From 73a9d53a37f6ce864b68dda1b07e92a0fed8c8ba Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Tue, 31 Mar 2026 01:57:32 -0700 Subject: [PATCH 01/83] CI: Add CodeQL workflow for GitHub Actions security scanning (#1408) * CI: Add CodeQL workflow for GitHub Actions security scanning * Update .github/workflows/codeql.yml --- .github/workflows/codeql.yml | 54 ++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 .github/workflows/codeql.yml diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml new file mode 100644 index 000000000..a9855cf48 --- /dev/null +++ b/.github/workflows/codeql.yml @@ -0,0 +1,54 @@ +# 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. +# + +name: "CodeQL" + +on: + push: + branches: [ "main" ] + pull_request: + branches: [ "main" ] + schedule: + - cron: '16 4 * * 1' + +permissions: + contents: read + +jobs: + analyze: + name: Analyze Actions + runs-on: ubuntu-latest + permissions: + contents: read + security-events: write + + steps: + - name: Checkout repository + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6 + with: + persist-credentials: false + + - name: Initialize CodeQL + uses: github/codeql-action/init@c793b717bc78562f491db7b0e93a3a178b099162 # v4 + with: + languages: actions + + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@c793b717bc78562f491db7b0e93a3a178b099162 # v4 + with: + category: "/language:actions" From 24994099e41a4e933f883557e2bce1a963bac0ea Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 31 Mar 2026 14:09:16 -0400 Subject: [PATCH 02/83] ci: update codespell paths (#1469) * Update path so it works well with pre-commit * Prefix path with asterisk so we get matching in both CI and pre-commit * Update paths for codespell --- pyproject.toml | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d05a64083..327199d1a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,12 +170,15 @@ extend-allowed-calls = ["datafusion.lit", "lit"] "docs/*" = ["D"] "docs/source/conf.py" = ["ANN001", "ERA001", "INP001"] +# CI and pre-commit invoke codespell with different paths, so we have a little +# redundancy here, and we intentionally drop python in the path. [tool.codespell] skip = [ - "./python/tests/test_functions.py", - "./target", + "*/tests/test_functions.py", + "*/target", + "./uv.lock", "uv.lock", - "./examples/tpch/answers_sf1/*", + "*/tpch/answers_sf1/*", ] count = true ignore-words-list = ["IST", "ans"] From 0113a6ee55cc61f9ebd897ae8cfc9213f560e468 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 2 Apr 2026 17:47:47 -0400 Subject: [PATCH 03/83] Add missing datetime functions (#1467) * Add missing datetime functions: make_time, current_timestamp, date_format Closes #1451. Adds make_time Rust binding and Python wrapper, and adds current_timestamp (alias for now) and date_format (alias for to_char) Python functions. Co-Authored-By: Claude Opus 4.6 (1M context) * Add unit tests for make_time, current_timestamp, and date_format Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- crates/core/src/functions.rs | 2 ++ python/datafusion/functions.py | 36 ++++++++++++++++++++++++++++++++++ python/tests/test_functions.py | 33 +++++++++++++++++++++++++++++++ 3 files changed, 71 insertions(+) diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index c32134054..6996dca94 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -616,6 +616,7 @@ expr_fn!(date_part, part date); expr_fn!(date_trunc, part date); expr_fn!(date_bin, stride source origin); expr_fn!(make_date, year month day); +expr_fn!(make_time, hour minute second); expr_fn!(to_char, datetime format); expr_fn!(translate, string from to, "Replaces each character in string that matches a character in the from set with the corresponding character in the to set. If from is longer than to, occurrences of the extra characters in from are deleted."); @@ -974,6 +975,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(date_part))?; m.add_wrapped(wrap_pyfunction!(date_trunc))?; m.add_wrapped(wrap_pyfunction!(make_date))?; + m.add_wrapped(wrap_pyfunction!(make_time))?; m.add_wrapped(wrap_pyfunction!(digest))?; m.add_wrapped(wrap_pyfunction!(ends_with))?; m.add_wrapped(wrap_pyfunction!(exp))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index f062cbfce..3c8d2bcee 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -128,7 +128,9 @@ "cume_dist", "current_date", "current_time", + "current_timestamp", "date_bin", + "date_format", "date_part", "date_trunc", "datepart", @@ -200,6 +202,7 @@ "make_array", "make_date", "make_list", + "make_time", "max", "md5", "mean", @@ -1948,6 +1951,15 @@ def now() -> Expr: return Expr(f.now()) +def current_timestamp() -> Expr: + """Returns the current timestamp in nanoseconds. + + See Also: + This is an alias for :py:func:`now`. + """ + return now() + + def to_char(arg: Expr, formatter: Expr) -> Expr: """Returns a string representation of a date, time, timestamp or duration. @@ -1970,6 +1982,15 @@ def to_char(arg: Expr, formatter: Expr) -> Expr: return Expr(f.to_char(arg.expr, formatter.expr)) +def date_format(arg: Expr, formatter: Expr) -> Expr: + """Returns a string representation of a date, time, timestamp or duration. + + See Also: + This is an alias for :py:func:`to_char`. + """ + return to_char(arg, formatter) + + def _unwrap_exprs(args: tuple[Expr, ...]) -> list: return [arg.expr for arg in args] @@ -2270,6 +2291,21 @@ def make_date(year: Expr, month: Expr, day: Expr) -> Expr: return Expr(f.make_date(year.expr, month.expr, day.expr)) +def make_time(hour: Expr, minute: Expr, second: Expr) -> Expr: + """Make a time from hour, minute and second component parts. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"h": [12], "m": [30], "s": [0]}) + >>> result = df.select( + ... dfn.functions.make_time(dfn.col("h"), dfn.col("m"), + ... dfn.col("s")).alias("t")) + >>> result.collect_column("t")[0].as_py() + datetime.time(12, 30) + """ + return Expr(f.make_time(hour.expr, minute.expr, second.expr)) + + def translate(string: Expr, from_val: Expr, to_val: Expr) -> Expr: """Replaces the characters in ``from_val`` with the counterpart in ``to_val``. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 37d349c58..08420826d 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -1107,6 +1107,39 @@ def test_today_alias_matches_current_date(df): assert result.column(0) == result.column(1) +def test_current_timestamp_alias_matches_now(df): + result = df.select( + f.now().alias("now"), + f.current_timestamp().alias("current_timestamp"), + ).collect()[0] + + assert result.column(0) == result.column(1) + + +def test_date_format_alias_matches_to_char(df): + result = df.select( + f.to_char( + f.to_timestamp(literal("2021-01-01T00:00:00")), literal("%Y/%m/%d") + ).alias("to_char"), + f.date_format( + f.to_timestamp(literal("2021-01-01T00:00:00")), literal("%Y/%m/%d") + ).alias("date_format"), + ).collect()[0] + + assert result.column(0) == result.column(1) + assert result.column(0)[0].as_py() == "2021/01/01" + + +def test_make_time(df): + ctx = SessionContext() + df_time = ctx.from_pydict({"h": [12], "m": [30], "s": [0]}) + result = df_time.select( + f.make_time(column("h"), column("m"), column("s")).alias("t") + ).collect()[0] + + assert result.column(0)[0].as_py() == time(12, 30) + + def test_arrow_cast(df): df = df.select( # we use `string_literal` to return utf8 instead of `literal` which returns From be8dd9d08fd284cf1747a2c1b965d9c95fff117c Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 3 Apr 2026 09:37:00 -0400 Subject: [PATCH 04/83] Add AI skill to check current repository against upstream APIs (#1460) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Initial commit for skill to check upstream repo * Add instructions on using the check-upstream skill * Add FFI type coverage and implementation pattern to check-upstream skill Document the full FFI type pipeline (Rust PyO3 wrapper → Protocol type → Python wrapper → ABC base class → exports → example) and catalog which upstream datafusion-ffi types are supported, which have been evaluated as not needing direct exposure, and how to check for new gaps. Co-Authored-By: Claude Opus 4.6 (1M context) * Update check-upstream skill to include FFI types as a checkable area Add "ffi types" to the argument-hint and description so users can invoke the skill with `/check-upstream ffi types`. Also add pipeline verification step to ensure each supported FFI type has the full end-to-end chain (PyO3 wrapper, Protocol, Python wrapper with type hints, ABC, exports). Co-Authored-By: Claude Opus 4.6 (1M context) * Move FFI Types section alongside other areas to check Section 7 (FFI Types) was incorrectly placed after the Output Format and Implementation Pattern sections. Move it to sit after Section 6 (SessionContext Methods), consistent with the other checkable areas. Co-Authored-By: Claude Opus 4.6 (1M context) * Replace static FFI type list with dynamic discovery instruction The supported FFI types list would go stale as new types are added. Replace it with a grep instruction to discover them at check time, keeping only the "evaluated and not requiring exposure" list which captures rationale not derivable from code. Co-Authored-By: Claude Opus 4.6 (1M context) * Make Python API the source of truth for upstream coverage checks Functions exposed in Python (e.g., as aliases of other Rust bindings) were being falsely reported as missing because they lacked a dedicated #[pyfunction] in Rust. The user-facing API is the Python layer, so coverage should be measured there. Co-Authored-By: Claude Opus 4.6 (1M context) * Add exclusion list for DataFrame methods already covered by Python API show_limit is covered by DataFrame.show() and with_param_values is covered by SessionContext.sql(param_values=...), so neither needs separate exposure. Co-Authored-By: Claude Opus 4.6 (1M context) * Move skills to .ai/skills/ for tool-agnostic discoverability Moves the canonical skill definitions from .claude/skills/ to .ai/skills/ and replaces .claude/skills with a symlink, so Claude Code still discovers them while other AI agents can find them in a tool-neutral location. Co-Authored-By: Claude Opus 4.6 (1M context) * Add AGENTS.md for tool-agnostic agent instructions with CLAUDE.md symlink AGENTS.md points agents to .ai/skills/ for skill discovery. CLAUDE.md symlinks to it so Claude Code picks it up as project instructions. Co-Authored-By: Claude Opus 4.6 (1M context) * Make README upstream coverage section tool-agnostic Remove Claude Code references and update skill path from .claude/skills/ to .ai/skills/ to match the new tool-neutral directory structure. Co-Authored-By: Claude Opus 4.6 (1M context) * Add GitHub issue lookup step to check-upstream skill When gaps are identified, search open issues at apache/datafusion-python before reporting. Existing issues are linked in the report rather than duplicated. Co-Authored-By: Claude Opus 4.6 (1M context) * Require Python test coverage in issues created by check-upstream skill Co-Authored-By: Claude Opus 4.6 (1M context) * Add license text --------- Co-authored-by: Claude Opus 4.6 (1M context) --- .ai/skills/check-upstream/SKILL.md | 382 +++++++++++++++++++++++++++++ .claude/skills | 1 + AGENTS.md | 27 ++ CLAUDE.md | 1 + README.md | 27 ++ 5 files changed, 438 insertions(+) create mode 100644 .ai/skills/check-upstream/SKILL.md create mode 120000 .claude/skills create mode 100644 AGENTS.md create mode 120000 CLAUDE.md diff --git a/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md new file mode 100644 index 000000000..f77210371 --- /dev/null +++ b/.ai/skills/check-upstream/SKILL.md @@ -0,0 +1,382 @@ + + +--- +name: check-upstream +description: Check if upstream Apache DataFusion features (functions, DataFrame ops, SessionContext methods, FFI types) are exposed in this Python project. Use when adding missing functions, auditing API coverage, or ensuring parity with upstream. +argument-hint: [area] (e.g., "scalar functions", "aggregate functions", "window functions", "dataframe", "session context", "ffi types", "all") +--- + +# Check Upstream DataFusion Feature Coverage + +You are auditing the datafusion-python project to find features from the upstream Apache DataFusion Rust library that are **not yet exposed** in this Python binding project. Your goal is to identify gaps and, if asked, implement the missing bindings. + +**IMPORTANT: The Python API is the source of truth for coverage.** A function or method is considered "exposed" if it exists in the Python API (e.g., `python/datafusion/functions.py`), even if there is no corresponding entry in the Rust bindings. Many upstream functions are aliases of other functions — the Python layer can expose these aliases by calling a different underlying Rust binding. Do NOT report a function as missing if it appears in the Python `__all__` list and has a working implementation, regardless of whether a matching `#[pyfunction]` exists in Rust. + +## Areas to Check + +The user may specify an area via `$ARGUMENTS`. If no area is specified or "all" is given, check all areas. + +### 1. Scalar Functions + +**Upstream source of truth:** +- Rust docs: https://docs.rs/datafusion/latest/datafusion/functions/index.html +- User docs: https://datafusion.apache.org/user-guide/sql/scalar_functions.html + +**Where they are exposed in this project:** +- Python API: `python/datafusion/functions.py` — each function wraps a call to `datafusion._internal.functions` +- Rust bindings: `crates/core/src/functions.rs` — `#[pyfunction]` definitions registered via `init_module()` + +**How to check:** +1. Fetch the upstream scalar function documentation page +2. Compare against functions listed in `python/datafusion/functions.py` (check the `__all__` list and function definitions) +3. A function is covered if it exists in the Python API — it does NOT need a dedicated Rust `#[pyfunction]`. Many functions are aliases that reuse another function's Rust binding. +4. Only report functions that are missing from the Python `__all__` list / function definitions + +### 2. Aggregate Functions + +**Upstream source of truth:** +- Rust docs: https://docs.rs/datafusion/latest/datafusion/functions_aggregate/index.html +- User docs: https://datafusion.apache.org/user-guide/sql/aggregate_functions.html + +**Where they are exposed in this project:** +- Python API: `python/datafusion/functions.py` (aggregate functions are mixed in with scalar functions) +- Rust bindings: `crates/core/src/functions.rs` + +**How to check:** +1. Fetch the upstream aggregate function documentation page +2. Compare against aggregate functions in `python/datafusion/functions.py` (check `__all__` list and function definitions) +3. A function is covered if it exists in the Python API, even if it aliases another function's Rust binding +4. Report only functions missing from the Python API + +### 3. Window Functions + +**Upstream source of truth:** +- Rust docs: https://docs.rs/datafusion/latest/datafusion/functions_window/index.html +- User docs: https://datafusion.apache.org/user-guide/sql/window_functions.html + +**Where they are exposed in this project:** +- Python API: `python/datafusion/functions.py` (window functions like `rank`, `dense_rank`, `lag`, `lead`, etc.) +- Rust bindings: `crates/core/src/functions.rs` + +**How to check:** +1. Fetch the upstream window function documentation page +2. Compare against window functions in `python/datafusion/functions.py` (check `__all__` list and function definitions) +3. A function is covered if it exists in the Python API, even if it aliases another function's Rust binding +4. Report only functions missing from the Python API + +### 4. Table Functions + +**Upstream source of truth:** +- Rust docs: https://docs.rs/datafusion/latest/datafusion/functions_table/index.html +- User docs: https://datafusion.apache.org/user-guide/sql/table_functions.html (if available) + +**Where they are exposed in this project:** +- Python API: `python/datafusion/functions.py` and `python/datafusion/user_defined.py` (TableFunction/udtf) +- Rust bindings: `crates/core/src/functions.rs` and `crates/core/src/udtf.rs` + +**How to check:** +1. Fetch the upstream table function documentation +2. Compare against what's available in the Python API +3. A function is covered if it exists in the Python API, even if it aliases another function's Rust binding +4. Report only functions missing from the Python API + +### 5. DataFrame Operations + +**Upstream source of truth:** +- Rust docs: https://docs.rs/datafusion/latest/datafusion/dataframe/struct.DataFrame.html + +**Where they are exposed in this project:** +- Python API: `python/datafusion/dataframe.py` — the `DataFrame` class +- Rust bindings: `crates/core/src/dataframe.rs` — `PyDataFrame` with `#[pymethods]` + +**Evaluated and not requiring separate Python exposure:** +- `show_limit` — already covered by `DataFrame.show()`, which provides the same functionality with a simpler API +- `with_param_values` — already covered by the `param_values` argument on `SessionContext.sql()`, which accomplishes the same thing more robustly + +**How to check:** +1. Fetch the upstream DataFrame documentation page listing all methods +2. Compare against methods in `python/datafusion/dataframe.py` — this is the source of truth for coverage +3. The Rust bindings (`crates/core/src/dataframe.rs`) may be consulted for context, but a method is covered if it exists in the Python API +4. Check against the "evaluated and not requiring exposure" list before flagging as a gap +5. Report only methods missing from the Python API + +### 6. SessionContext Methods + +**Upstream source of truth:** +- Rust docs: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html + +**Where they are exposed in this project:** +- Python API: `python/datafusion/context.py` — the `SessionContext` class +- Rust bindings: `crates/core/src/context.rs` — `PySessionContext` with `#[pymethods]` + +**How to check:** +1. Fetch the upstream SessionContext documentation page listing all methods +2. Compare against methods in `python/datafusion/context.py` — this is the source of truth for coverage +3. The Rust bindings (`crates/core/src/context.rs`) may be consulted for context, but a method is covered if it exists in the Python API +4. Report only methods missing from the Python API + +### 7. FFI Types (datafusion-ffi) + +**Upstream source of truth:** +- Crate source: https://github.com/apache/datafusion/tree/main/datafusion/ffi/src +- Rust docs: https://docs.rs/datafusion-ffi/latest/datafusion_ffi/ + +**Where they are exposed in this project:** +- Rust bindings: various files under `crates/core/src/` and `crates/util/src/` +- FFI example: `examples/datafusion-ffi-example/src/` +- Dependency declared in root `Cargo.toml` and `crates/core/Cargo.toml` + +**Discovering currently supported FFI types:** +Grep for `use datafusion_ffi::` in `crates/core/src/` and `crates/util/src/` to find all FFI types currently imported and used. + +**Evaluated and not requiring direct Python exposure:** +These upstream FFI types have been reviewed and do not need to be independently exposed to end users: +- `FFI_ExecutionPlan` — already used indirectly through table providers; no need for direct exposure +- `FFI_PhysicalExpr` / `FFI_PhysicalSortExpr` — internal physical planning types not expected to be needed by end users +- `FFI_RecordBatchStream` — one level deeper than FFI_ExecutionPlan, used internally when execution plans stream results +- `FFI_SessionRef` / `ForeignSession` — session sharing across FFI; Python manages sessions natively via SessionContext +- `FFI_SessionConfig` — Python can configure sessions natively without FFI +- `FFI_ConfigOptions` / `FFI_TableOptions` — internal configuration plumbing +- `FFI_PlanProperties` / `FFI_Boundedness` / `FFI_EmissionType` — read from existing plans, not user-facing +- `FFI_Partitioning` — supporting type for physical planning +- Supporting/utility types (`FFI_Option`, `FFI_Result`, `WrappedSchema`, `WrappedArray`, `FFI_ColumnarValue`, `FFI_Volatility`, `FFI_InsertOp`, `FFI_AccumulatorArgs`, `FFI_Accumulator`, `FFI_GroupsAccumulator`, `FFI_EmitTo`, `FFI_AggregateOrderSensitivity`, `FFI_PartitionEvaluator`, `FFI_PartitionEvaluatorArgs`, `FFI_Range`, `FFI_SortOptions`, `FFI_Distribution`, `FFI_ExprProperties`, `FFI_SortProperties`, `FFI_Interval`, `FFI_TableProviderFilterPushDown`, `FFI_TableType`) — used as building blocks within the types above, not independently exposed + +**How to check:** +1. Discover currently supported types by grepping for `use datafusion_ffi::` in `crates/core/src/` and `crates/util/src/`, then compare against the upstream `datafusion-ffi` crate's `lib.rs` exports +2. If new FFI types appear upstream, evaluate whether they represent a user-facing capability +3. Check against the "evaluated and not requiring exposure" list before flagging as a gap +4. Report any genuinely new types that enable user-facing functionality +5. For each currently supported FFI type, verify the full pipeline is present using the checklist from "Adding a New FFI Type": + - Rust PyO3 wrapper with `from_pycapsule()` method + - Python Protocol type (e.g., `ScalarUDFExportable`) for FFI objects + - Python wrapper class with full type hints on all public methods + - ABC base class (if the type can be user-implemented) + - Registered in Rust `init_module()` and Python `__init__.py` + - FFI example in `examples/datafusion-ffi-example/` + - Type appears in union type hints where accepted + +## Checking for Existing GitHub Issues + +After identifying missing APIs, search the open issues at https://github.com/apache/datafusion-python/issues for each gap to see if an issue already exists requesting that API be exposed. Search using the function or method name as the query. + +- If an existing issue is found, include a link to it in the report. Do NOT create a new issue. +- If no existing issue is found, note that no issue exists yet. If the user asks to create issues for missing APIs, each issue should specify that Python test coverage is required as part of the implementation. + +## Output Format + +For each area checked, produce a report like: + +``` +## [Area Name] Coverage Report + +### Currently Exposed (X functions/methods) +- list of what's already available + +### Missing from Upstream (Y functions/methods) +- function_name — brief description of what it does (existing issue: #123) +- function_name — brief description of what it does (no existing issue) + +### Notes +- Any relevant observations about partial implementations, naming differences, etc. +``` + +## Implementation Pattern + +If the user asks you to implement missing features, follow these patterns: + +### Adding a New Function (Scalar/Aggregate/Window) + +**Step 1: Rust binding** in `crates/core/src/functions.rs`: +```rust +#[pyfunction] +#[pyo3(signature = (arg1, arg2))] +fn new_function_name(arg1: PyExpr, arg2: PyExpr) -> PyResult { + Ok(datafusion::functions::module::expr_fn::new_function_name(arg1.expr, arg2.expr).into()) +} +``` +Then register in `init_module()`: +```rust +m.add_wrapped(wrap_pyfunction!(new_function_name))?; +``` + +**Step 2: Python wrapper** in `python/datafusion/functions.py`: +```python +def new_function_name(arg1: Expr, arg2: Expr) -> Expr: + """Description of what the function does. + + Args: + arg1: Description of first argument. + arg2: Description of second argument. + + Returns: + Description of return value. + """ + return Expr(f.new_function_name(arg1.expr, arg2.expr)) +``` +Add to `__all__` list. + +### Adding a New DataFrame Method + +**Step 1: Rust binding** in `crates/core/src/dataframe.rs`: +```rust +#[pymethods] +impl PyDataFrame { + fn new_method(&self, py: Python, param: PyExpr) -> PyDataFusionResult { + let df = self.df.as_ref().clone().new_method(param.into())?; + Ok(Self::new(df)) + } +} +``` + +**Step 2: Python wrapper** in `python/datafusion/dataframe.py`: +```python +def new_method(self, param: Expr) -> DataFrame: + """Description of the method.""" + return DataFrame(self.df.new_method(param.expr)) +``` + +### Adding a New SessionContext Method + +**Step 1: Rust binding** in `crates/core/src/context.rs`: +```rust +#[pymethods] +impl PySessionContext { + pub fn new_method(&self, py: Python, param: String) -> PyDataFusionResult { + let df = wait_for_future(py, self.ctx.new_method(¶m))?; + Ok(PyDataFrame::new(df)) + } +} +``` + +**Step 2: Python wrapper** in `python/datafusion/context.py`: +```python +def new_method(self, param: str) -> DataFrame: + """Description of the method.""" + return DataFrame(self.ctx.new_method(param)) +``` + +### Adding a New FFI Type + +FFI types require a full pipeline from C struct through to a typed Python wrapper. Each layer must be present. + +**Step 1: Rust PyO3 wrapper class** in a new or existing file under `crates/core/src/`: +```rust +use datafusion_ffi::new_type::FFI_NewType; + +#[pyclass(from_py_object, frozen, name = "RawNewType", module = "datafusion.module_name", subclass)] +pub struct PyNewType { + pub inner: Arc, +} + +#[pymethods] +impl PyNewType { + #[staticmethod] + fn from_pycapsule(obj: &Bound<'_, PyAny>) -> PyDataFusionResult { + let capsule = obj + .getattr("__datafusion_new_type__")? + .call0()? + .downcast::()?; + let ffi_ptr = unsafe { capsule.reference::() }; + let provider: Arc = ffi_ptr.into(); + Ok(Self { inner: provider }) + } + + fn some_method(&self) -> PyResult<...> { + // wrap inner trait method + } +} +``` +Register in the appropriate `init_module()`: +```rust +m.add_class::()?; +``` + +**Step 2: Python Protocol type** in the appropriate Python module (e.g., `python/datafusion/catalog.py`): +```python +class NewTypeExportable(Protocol): + """Type hint for objects providing a __datafusion_new_type__ PyCapsule.""" + + def __datafusion_new_type__(self) -> object: ... +``` + +**Step 3: Python wrapper class** in the same module: +```python +class NewType: + """Description of the type. + + This class wraps a DataFusion NewType, which can be created from a native + Python implementation or imported from an FFI-compatible library. + """ + + def __init__( + self, + new_type: df_internal.module_name.RawNewType | NewTypeExportable, + ) -> None: + if isinstance(new_type, df_internal.module_name.RawNewType): + self._raw = new_type + else: + self._raw = df_internal.module_name.RawNewType.from_pycapsule(new_type) + + def some_method(self) -> ReturnType: + """Description of the method.""" + return self._raw.some_method() +``` + +**Step 4: ABC base class** (if users should be able to subclass and provide custom implementations in Python): +```python +from abc import ABC, abstractmethod + +class NewTypeProvider(ABC): + """Abstract base class for implementing a custom NewType in Python.""" + + @abstractmethod + def some_method(self) -> ReturnType: + """Description of the method.""" + ... +``` + +**Step 5: Module exports** — add to the appropriate `__init__.py`: +- Add the wrapper class (`NewType`) to `python/datafusion/__init__.py` +- Add the ABC (`NewTypeProvider`) if applicable +- Add the Protocol type (`NewTypeExportable`) if it should be public + +**Step 6: FFI example** — add an example implementation under `examples/datafusion-ffi-example/src/`: +```rust +// examples/datafusion-ffi-example/src/new_type.rs +use datafusion_ffi::new_type::FFI_NewType; +// ... example showing how an external Rust library exposes this type via PyCapsule +``` + +**Checklist for each FFI type:** +- [ ] Rust PyO3 wrapper with `from_pycapsule()` method +- [ ] Python Protocol type (e.g., `NewTypeExportable`) for FFI objects +- [ ] Python wrapper class with full type hints on all public methods +- [ ] ABC base class (if the type can be user-implemented) +- [ ] Registered in Rust `init_module()` and Python `__init__.py` +- [ ] FFI example in `examples/datafusion-ffi-example/` +- [ ] Type appears in union type hints where accepted (e.g., `Table | TableProviderExportable`) + +## Important Notes + +- The upstream DataFusion version used by this project is specified in `crates/core/Cargo.toml` — check the `datafusion` dependency version to ensure you're comparing against the right upstream version. +- Some upstream features may intentionally not be exposed (e.g., internal-only APIs). Use judgment about what's user-facing. +- When fetching upstream docs, prefer the published docs.rs documentation as it matches the crate version. +- Function aliases (e.g., `array_append` / `list_append`) should both be exposed if upstream supports them. +- Check the `__all__` list in `functions.py` to see what's publicly exported vs just defined. diff --git a/.claude/skills b/.claude/skills new file mode 120000 index 000000000..6838a1160 --- /dev/null +++ b/.claude/skills @@ -0,0 +1 @@ +../.ai/skills \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..1853a84cd --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,27 @@ + + +# Agent Instructions + +This project uses AI agent skills stored in `.ai/skills/`. Each skill is a directory containing a `SKILL.md` file with instructions for performing a specific task. + +Skills follow the [Agent Skills](https://agentskills.io) open standard. Each skill directory contains: + +- `SKILL.md` — The skill definition with YAML frontmatter (name, description, argument-hint) and detailed instructions. +- Additional supporting files as needed. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 120000 index 000000000..47dc3e3d8 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md \ No newline at end of file diff --git a/README.md b/README.md index c24257876..7c1c71281 100644 --- a/README.md +++ b/README.md @@ -312,6 +312,33 @@ There are scripts in `ci/scripts` for running Rust and Python linters. ./ci/scripts/rust_toml_fmt.sh ``` +## Checking Upstream DataFusion Coverage + +This project includes an [AI agent skill](.ai/skills/check-upstream/SKILL.md) for auditing which +features from the upstream Apache DataFusion Rust library are not yet exposed in these Python +bindings. This is useful when adding missing functions, auditing API coverage, or ensuring parity +with upstream. + +The skill accepts an optional area argument: + +``` +scalar functions +aggregate functions +window functions +dataframe +session context +ffi types +all +``` + +If no argument is provided, it defaults to checking all areas. The skill will fetch the upstream +DataFusion documentation, compare it against the functions and methods exposed in this project, and +produce a coverage report listing what is currently exposed and what is missing. + +The skill definition lives in `.ai/skills/check-upstream/SKILL.md` and follows the +[Agent Skills](https://agentskills.io) open standard. It can be used by any AI coding agent that +supports skill discovery, or followed manually. + ## How to update dependencies To change test dependencies, change the `pyproject.toml` and run From 645d261ce3bc0b3b610c8d82422042b3e573e793 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 3 Apr 2026 13:51:43 -0400 Subject: [PATCH 05/83] Add missing string function `contains` (#1465) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add missing `contains` string function Expose the upstream DataFusion `contains(string, search_str)` function which returns true if search_str is found within string (case-sensitive). Note: the other functions from #1450 (instr, position, substring_index) already exist — instr and position are aliases for strpos, and substring_index is exposed as substr_index. Closes #1450 Co-Authored-By: Claude Opus 4.6 (1M context) * Add unit test for contains string function Co-Authored-By: Claude Opus 4.6 (1M context) * Update python/datafusion/functions.py Co-authored-by: Nuno Faria --------- Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Nuno Faria --- crates/core/src/functions.rs | 6 ++++++ python/datafusion/functions.py | 15 +++++++++++++++ python/tests/test_functions.py | 1 + 3 files changed, 22 insertions(+) diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 6996dca94..fefe14b3e 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -494,6 +494,11 @@ expr_fn!(length, string); expr_fn!(char_length, string); expr_fn!(chr, arg, "Returns the character with the given code."); expr_fn_vec!(coalesce); +expr_fn!( + contains, + string search_str, + "Return true if search_str is found within string (case-sensitive)." +); expr_fn!(cos, num); expr_fn!(cosh, num); expr_fn!(cot, num); @@ -961,6 +966,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(col))?; m.add_wrapped(wrap_pyfunction!(concat_ws))?; m.add_wrapped(wrap_pyfunction!(concat))?; + m.add_wrapped(wrap_pyfunction!(contains))?; m.add_wrapped(wrap_pyfunction!(corr))?; m.add_wrapped(wrap_pyfunction!(cos))?; m.add_wrapped(wrap_pyfunction!(cosh))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 3c8d2bcee..2ef2f0473 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -116,6 +116,7 @@ "col", "concat", "concat_ws", + "contains", "corr", "cos", "cosh", @@ -439,6 +440,20 @@ def digest(value: Expr, method: Expr) -> Expr: return Expr(f.digest(value.expr, method.expr)) +def contains(string: Expr, search_str: Expr) -> Expr: + """Returns true if ``search_str`` is found within ``string`` (case-sensitive). + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["the quick brown fox"]}) + >>> result = df.select( + ... dfn.functions.contains(dfn.col("a"), dfn.lit("brown")).alias("c")) + >>> result.collect_column("c")[0].as_py() + True + """ + return Expr(f.contains(string.expr, search_str.expr)) + + def concat(*args: Expr) -> Expr: """Concatenates the text representations of all the arguments. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 08420826d..db141fbe0 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -745,6 +745,7 @@ def test_array_function_obj_tests(stmt, py_expr): f.split_part(column("a"), literal("l"), literal(1)), pa.array(["He", "Wor", "!"]), ), + (f.contains(column("a"), literal("ell")), pa.array([True, False, False])), (f.starts_with(column("a"), literal("Wor")), pa.array([False, True, False])), (f.strpos(column("a"), literal("o")), pa.array([5, 2, 0], type=pa.int32())), ( From 0b6ea95a3d304a774bbe512bb70fbca332aa5426 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 3 Apr 2026 15:43:28 -0400 Subject: [PATCH 06/83] Add missing conditional functions (#1464) * Add missing conditional functions: greatest, least, nvl2, ifnull (#1449) Expose four conditional functions from upstream DataFusion that were not yet available in the Python bindings. Co-Authored-By: Claude Opus 4.6 (1M context) * Add unit tests for greatest, least, nvl2, and ifnull functions Tests cover multiple data types (integers, strings), null handling (all-null, partial-null), multiple arguments, and ifnull/nvl equivalence. Co-Authored-By: Claude Opus 4.6 (1M context) * Use standard alias docstring pattern for ifnull Co-Authored-By: Claude Opus 4.6 (1M context) * remove unused df fixture and fix parameter shadowing * Refactor conditional function tests into parametrized test suite Replace separate test functions for coalesce, greatest, least, nvl, nvl2, ifnull with a single parametrized test using a shared fixture. Adds coverage for nvl, nullif (previously untested), datetime and boolean types, literal fallbacks, and variadic calls. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- crates/core/src/functions.rs | 10 ++ python/datafusion/functions.py | 69 ++++++++ python/tests/test_functions.py | 289 +++++++++++++++++++++++++++------ 3 files changed, 319 insertions(+), 49 deletions(-) diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index fefe14b3e..3f07da95b 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -494,6 +494,8 @@ expr_fn!(length, string); expr_fn!(char_length, string); expr_fn!(chr, arg, "Returns the character with the given code."); expr_fn_vec!(coalesce); +expr_fn_vec!(greatest); +expr_fn_vec!(least); expr_fn!( contains, string search_str, @@ -548,6 +550,11 @@ expr_fn!( x y, "Returns x if x is not NULL otherwise returns y." ); +expr_fn!( + nvl2, + x y z, + "Returns y if x is not NULL; otherwise returns z." +); expr_fn!(nullif, arg_1 arg_2); expr_fn!( octet_length, @@ -989,6 +996,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(floor))?; m.add_wrapped(wrap_pyfunction!(from_unixtime))?; m.add_wrapped(wrap_pyfunction!(gcd))?; + m.add_wrapped(wrap_pyfunction!(greatest))?; // m.add_wrapped(wrap_pyfunction!(grouping))?; m.add_wrapped(wrap_pyfunction!(in_list))?; m.add_wrapped(wrap_pyfunction!(initcap))?; @@ -996,6 +1004,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(iszero))?; m.add_wrapped(wrap_pyfunction!(levenshtein))?; m.add_wrapped(wrap_pyfunction!(lcm))?; + m.add_wrapped(wrap_pyfunction!(least))?; m.add_wrapped(wrap_pyfunction!(left))?; m.add_wrapped(wrap_pyfunction!(length))?; m.add_wrapped(wrap_pyfunction!(ln))?; @@ -1013,6 +1022,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(named_struct))?; m.add_wrapped(wrap_pyfunction!(nanvl))?; m.add_wrapped(wrap_pyfunction!(nvl))?; + m.add_wrapped(wrap_pyfunction!(nvl2))?; m.add_wrapped(wrap_pyfunction!(now))?; m.add_wrapped(wrap_pyfunction!(nullif))?; m.add_wrapped(wrap_pyfunction!(octet_length))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 2ef2f0473..f1ea3d256 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -152,6 +152,8 @@ "floor", "from_unixtime", "gcd", + "greatest", + "ifnull", "in_list", "initcap", "isnan", @@ -160,6 +162,7 @@ "last_value", "lcm", "lead", + "least", "left", "length", "levenshtein", @@ -216,6 +219,7 @@ "ntile", "nullif", "nvl", + "nvl2", "octet_length", "order_by", "overlay", @@ -1045,6 +1049,34 @@ def gcd(x: Expr, y: Expr) -> Expr: return Expr(f.gcd(x.expr, y.expr)) +def greatest(*args: Expr) -> Expr: + """Returns the greatest value from a list of expressions. + + Returns NULL if all expressions are NULL. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 3], "b": [2, 1]}) + >>> result = df.select( + ... dfn.functions.greatest(dfn.col("a"), dfn.col("b")).alias("greatest")) + >>> result.collect_column("greatest")[0].as_py() + 2 + >>> result.collect_column("greatest")[1].as_py() + 3 + """ + exprs = [arg.expr for arg in args] + return Expr(f.greatest(*exprs)) + + +def ifnull(x: Expr, y: Expr) -> Expr: + """Returns ``x`` if ``x`` is not NULL. Otherwise returns ``y``. + + See Also: + This is an alias for :py:func:`nvl`. + """ + return nvl(x, y) + + def initcap(string: Expr) -> Expr: """Set the initial letter of each word to capital. @@ -1098,6 +1130,25 @@ def lcm(x: Expr, y: Expr) -> Expr: return Expr(f.lcm(x.expr, y.expr)) +def least(*args: Expr) -> Expr: + """Returns the least value from a list of expressions. + + Returns NULL if all expressions are NULL. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 3], "b": [2, 1]}) + >>> result = df.select( + ... dfn.functions.least(dfn.col("a"), dfn.col("b")).alias("least")) + >>> result.collect_column("least")[0].as_py() + 1 + >>> result.collect_column("least")[1].as_py() + 1 + """ + exprs = [arg.expr for arg in args] + return Expr(f.least(*exprs)) + + def left(string: Expr, n: Expr) -> Expr: """Returns the first ``n`` characters in the ``string``. @@ -1282,6 +1333,24 @@ def nvl(x: Expr, y: Expr) -> Expr: return Expr(f.nvl(x.expr, y.expr)) +def nvl2(x: Expr, y: Expr, z: Expr) -> Expr: + """Returns ``y`` if ``x`` is not NULL. Otherwise returns ``z``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [None, 1], "b": [10, 20], "c": [30, 40]}) + >>> result = df.select( + ... dfn.functions.nvl2( + ... dfn.col("a"), dfn.col("b"), dfn.col("c")).alias("nvl2") + ... ) + >>> result.collect_column("nvl2")[0].as_py() + 30 + >>> result.collect_column("nvl2")[1].as_py() + 20 + """ + return Expr(f.nvl2(x.expr, y.expr, z.expr)) + + def octet_length(arg: Expr) -> Expr: """Returns the number of bytes of a string. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index db141fbe0..74fcbffb4 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -1410,62 +1410,253 @@ def test_alias_with_metadata(df): assert df.schema().field("b").metadata == {b"key": b"value"} -def test_coalesce(df): - # Create a DataFrame with null values +@pytest.fixture +def df_with_nulls(): ctx = SessionContext() + # Rows: + # 0: both values present + # 1: a/d/h/k null, b/e/i/l present + # 2: a/d/h/k present, b/e/i/l null + # 3: all null batch = pa.RecordBatch.from_arrays( [ - pa.array(["Hello", None, "!"]), # string column with null - pa.array([4, None, 6]), # integer column with null - pa.array(["hello ", None, " !"]), # string column with null + pa.array([1, None, 3, None], type=pa.int64()), + pa.array([5, 10, None, None], type=pa.int64()), + pa.array([20, 30, 40, None], type=pa.int64()), + pa.array(["apple", None, "cherry", None], type=pa.utf8()), + pa.array(["banana", "date", None, None], type=pa.utf8()), + pa.array(["x", "y", "z", None], type=pa.utf8()), pa.array( [ - datetime(2022, 12, 31, tzinfo=DEFAULT_TZ), + datetime(2020, 1, 1, tzinfo=DEFAULT_TZ), None, - datetime(2020, 7, 2, tzinfo=DEFAULT_TZ), - ] - ), # datetime with null - pa.array([False, None, True]), # boolean column with null + datetime(2025, 6, 15, tzinfo=DEFAULT_TZ), + None, + ], + type=pa.timestamp("us", tz="UTC"), + ), + pa.array( + [ + datetime(2022, 7, 4, tzinfo=DEFAULT_TZ), + datetime(2023, 12, 25, tzinfo=DEFAULT_TZ), + None, + None, + ], + type=pa.timestamp("us", tz="UTC"), + ), + pa.array([True, None, False, None], type=pa.bool_()), + pa.array([False, True, None, None], type=pa.bool_()), ], - names=["a", "b", "c", "d", "e"], - ) - df_with_nulls = ctx.create_dataframe([[batch]]) - - # Test coalesce with different data types - result_df = df_with_nulls.select( - f.coalesce(column("a"), literal("default")).alias("a_coalesced"), - f.coalesce(column("b"), literal(0)).alias("b_coalesced"), - f.coalesce(column("c"), literal("default")).alias("c_coalesced"), - f.coalesce(column("d"), literal(datetime(2000, 1, 1, tzinfo=DEFAULT_TZ))).alias( - "d_coalesced" - ), - f.coalesce(column("e"), literal(value=False)).alias("e_coalesced"), + names=["a", "b", "c", "d", "e", "g", "h", "i", "k", "l"], ) + return ctx.create_dataframe([[batch]]) - result = result_df.collect()[0] - # Verify results - assert result.column(0) == pa.array( - ["Hello", "default", "!"], type=pa.string_view() - ) - assert result.column(1) == pa.array([4, 0, 6], type=pa.int64()) - assert result.column(2) == pa.array( - ["hello ", "default", " !"], type=pa.string_view() - ) - assert result.column(3).to_pylist() == [ - datetime(2022, 12, 31, tzinfo=DEFAULT_TZ), - datetime(2000, 1, 1, tzinfo=DEFAULT_TZ), - datetime(2020, 7, 2, tzinfo=DEFAULT_TZ), - ] - assert result.column(4) == pa.array([False, False, True], type=pa.bool_()) - - # Test multiple arguments - result_df = df_with_nulls.select( - f.coalesce(column("a"), literal(None), literal("fallback")).alias( - "multi_coalesce" - ) - ) - result = result_df.collect()[0] - assert result.column(0) == pa.array( - ["Hello", "fallback", "!"], type=pa.string_view() - ) +@pytest.mark.parametrize( + ("expr", "expected"), + [ + pytest.param( + f.greatest(column("a"), column("b")), + pa.array([5, 10, 3, None], type=pa.int64()), + id="greatest_int", + ), + pytest.param( + f.greatest(column("d"), column("e")), + pa.array(["banana", "date", "cherry", None], type=pa.utf8()), + id="greatest_str", + ), + pytest.param( + f.least(column("a"), column("b")), + pa.array([1, 10, 3, None], type=pa.int64()), + id="least_int", + ), + pytest.param( + f.least(column("d"), column("e")), + pa.array(["apple", "date", "cherry", None], type=pa.utf8()), + id="least_str", + ), + pytest.param( + f.coalesce(column("a"), column("b"), column("c")), + pa.array([1, 10, 3, None], type=pa.int64()), + id="coalesce_int", + ), + pytest.param( + f.coalesce(column("d"), column("e"), column("g")), + pa.array(["apple", "date", "cherry", None], type=pa.utf8()), + id="coalesce_str", + ), + pytest.param( + f.nvl(column("a"), column("c")), + pa.array([1, 30, 3, None], type=pa.int64()), + id="nvl_int", + ), + pytest.param( + f.nvl(column("d"), column("g")), + pa.array(["apple", "y", "cherry", None], type=pa.utf8()), + id="nvl_str", + ), + pytest.param( + f.ifnull(column("a"), column("c")), + pa.array([1, 30, 3, None], type=pa.int64()), + id="ifnull_int", + ), + pytest.param( + f.ifnull(column("d"), column("g")), + pa.array(["apple", "y", "cherry", None], type=pa.utf8()), + id="ifnull_str", + ), + pytest.param( + f.nvl2(column("a"), column("b"), column("c")), + pa.array([5, 30, None, None], type=pa.int64()), + id="nvl2_int", + ), + pytest.param( + f.nvl2(column("d"), column("e"), column("g")), + pa.array(["banana", "y", None, None], type=pa.utf8()), + id="nvl2_str", + ), + pytest.param( + f.nullif(column("a"), column("b")), + pa.array([1, None, 3, None], type=pa.int64()), + id="nullif_int", + ), + pytest.param( + f.nullif(column("d"), column("e")), + pa.array(["apple", None, "cherry", None], type=pa.utf8()), + id="nullif_str", + ), + pytest.param( + f.nullif(column("a"), literal(1)), + pa.array([None, None, 3, None], type=pa.int64()), + id="nullif_equal_values", + ), + pytest.param( + f.greatest(column("a"), column("b"), column("c")), + pa.array([20, 30, 40, None], type=pa.int64()), + id="greatest_variadic", + ), + pytest.param( + f.least(column("a"), column("b"), column("c")), + pa.array([1, 10, 3, None], type=pa.int64()), + id="least_variadic", + ), + pytest.param( + f.greatest(column("a"), literal(2)), + pa.array([2, 2, 3, 2], type=pa.int64()), + id="greatest_literal", + ), + pytest.param( + f.least(column("a"), literal(2)), + pa.array([1, 2, 2, 2], type=pa.int64()), + id="least_literal", + ), + pytest.param( + f.coalesce(column("a"), literal(0)), + pa.array([1, 0, 3, 0], type=pa.int64()), + id="coalesce_literal_int", + ), + pytest.param( + f.coalesce(column("d"), literal("default")), + pa.array(["apple", "default", "cherry", "default"], type=pa.string_view()), + id="coalesce_literal_str", + ), + pytest.param( + f.nvl(column("a"), literal(99)), + pa.array([1, 99, 3, 99], type=pa.int64()), + id="nvl_literal", + ), + pytest.param( + f.ifnull(column("d"), literal("unknown")), + pa.array(["apple", "unknown", "cherry", "unknown"], type=pa.string_view()), + id="ifnull_literal", + ), + pytest.param( + f.nvl2(column("a"), literal(1), literal(0)), + pa.array([1, 0, 1, 0], type=pa.int64()), + id="nvl2_literal", + ), + pytest.param( + f.greatest(column("h"), column("i")), + pa.array( + [ + datetime(2022, 7, 4, tzinfo=DEFAULT_TZ), + datetime(2023, 12, 25, tzinfo=DEFAULT_TZ), + datetime(2025, 6, 15, tzinfo=DEFAULT_TZ), + None, + ], + type=pa.timestamp("us", tz="UTC"), + ), + id="greatest_datetime", + ), + pytest.param( + f.least(column("h"), column("i")), + pa.array( + [ + datetime(2020, 1, 1, tzinfo=DEFAULT_TZ), + datetime(2023, 12, 25, tzinfo=DEFAULT_TZ), + datetime(2025, 6, 15, tzinfo=DEFAULT_TZ), + None, + ], + type=pa.timestamp("us", tz="UTC"), + ), + id="least_datetime", + ), + pytest.param( + f.coalesce(column("h"), column("i")), + pa.array( + [ + datetime(2020, 1, 1, tzinfo=DEFAULT_TZ), + datetime(2023, 12, 25, tzinfo=DEFAULT_TZ), + datetime(2025, 6, 15, tzinfo=DEFAULT_TZ), + None, + ], + type=pa.timestamp("us", tz="UTC"), + ), + id="coalesce_datetime", + ), + pytest.param( + f.nvl(column("k"), column("l")), + pa.array([True, True, False, None], type=pa.bool_()), + id="nvl_bool", + ), + pytest.param( + f.coalesce(column("k"), column("l")), + pa.array([True, True, False, None], type=pa.bool_()), + id="coalesce_bool", + ), + pytest.param( + f.nvl2(column("k"), column("k"), column("l")), + pa.array([True, True, False, None], type=pa.bool_()), + id="nvl2_bool", + ), + pytest.param( + f.coalesce( + column("h"), + literal(datetime(2000, 1, 1, tzinfo=DEFAULT_TZ)), + ), + pa.array( + [ + datetime(2020, 1, 1, tzinfo=DEFAULT_TZ), + datetime(2000, 1, 1, tzinfo=DEFAULT_TZ), + datetime(2025, 6, 15, tzinfo=DEFAULT_TZ), + datetime(2000, 1, 1, tzinfo=DEFAULT_TZ), + ], + type=pa.timestamp("us", tz="UTC"), + ), + id="coalesce_literal_datetime", + ), + pytest.param( + f.coalesce(column("k"), literal(value=False)), + pa.array([True, False, False, False], type=pa.bool_()), + id="coalesce_literal_bool", + ), + pytest.param( + f.coalesce(column("a"), literal(None), literal(99)), + pa.array([1, 99, 3, 99], type=pa.int64()), + id="coalesce_skip_null_literal", + ), + ], +) +def test_conditional_functions(df_with_nulls, expr, expected): + result = df_with_nulls.select(expr.alias("result")).collect()[0] + assert result.column(0) == expected From 16feeb136737ae45fac39f7a82cca2d88fd6224b Mon Sep 17 00:00:00 2001 From: Kevin Liu Date: Fri, 3 Apr 2026 12:47:31 -0700 Subject: [PATCH 07/83] Reduce peak memory usage during release builds to fix OOM on manylinux runners (#1445) * adjust swap to 8gb * modify profile.release --- .github/workflows/build.yml | 15 ++++++++++++++- Cargo.toml | 4 ++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 4b37046ee..7682d6cb0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -159,6 +159,19 @@ jobs: with: enable-cache: true + - name: Add extra swap for release build + if: inputs.build_mode == 'release' + run: | + set -euxo pipefail + sudo swapoff -a || true + sudo rm -f /swapfile + sudo fallocate -l 8G /swapfile || sudo dd if=/dev/zero of=/swapfile bs=1M count=8192 + sudo chmod 600 /swapfile + sudo mkswap /swapfile + sudo swapon /swapfile + free -h + swapon --show + - name: Build (release mode) uses: PyO3/maturin-action@v1 if: inputs.build_mode == 'release' @@ -233,7 +246,7 @@ jobs: set -euxo pipefail sudo swapoff -a || true sudo rm -f /swapfile - sudo fallocate -l 16G /swapfile || sudo dd if=/dev/zero of=/swapfile bs=1M count=16384 + sudo fallocate -l 8G /swapfile || sudo dd if=/dev/zero of=/swapfile bs=1M count=8192 sudo chmod 600 /swapfile sudo mkswap /swapfile sudo swapon /swapfile diff --git a/Cargo.toml b/Cargo.toml index 346f6da3e..3a34e204c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -64,8 +64,8 @@ pyo3-build-config = "0.28" datafusion-python-util = { path = "crates/util" } [profile.release] -lto = true -codegen-units = 1 +lto = "thin" +codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. From 8a35caea9ed01492742738f161fa5b4459d69402 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 4 Apr 2026 12:20:31 -0400 Subject: [PATCH 08/83] Add missing map functions (#1461) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add map functions (make_map, map_keys, map_values, map_extract, map_entries, element_at) Closes #1448 Co-Authored-By: Claude Opus 4.6 (1M context) * Add unit tests for map functions Co-Authored-By: Claude Opus 4.6 (1M context) * Remove redundant pyo3 element_at function element_at is already a Python-only alias for map_extract, so the Rust binding is unnecessary. Co-Authored-By: Claude Opus 4.6 (1M context) * Change make_map to accept a Python dictionary make_map now takes a dict for the common case and also supports separate keys/values lists for column expressions. Non-Expr keys and values are automatically converted to literals. Co-Authored-By: Claude Opus 4.6 (1M context) * Make map the primary function with make_map as alias map() now supports three calling conventions matching upstream: - map({"a": 1, "b": 2}) — from a Python dictionary - map([keys], [values]) — two lists that get zipped - map(k1, v1, k2, v2, ...) — variadic key-value pairs Non-Expr keys and values are automatically converted to literals. Co-Authored-By: Claude Opus 4.6 (1M context) * Improve map function docstrings - Add examples for all three map() calling conventions - Use clearer descriptions instead of jargon (no "zipped" or "variadic") - Break map_keys/map_values/map_extract/map_entries examples into two steps: create the map column first, then call the function Co-Authored-By: Claude Opus 4.6 (1M context) * Remove map() in favor of make_map(), fix docstrings, add validation - Remove map() function that shadowed Python builtin; make_map() is now the sole entry point for creating map expressions - Fix map_extract/element_at docstrings: missing keys return [None], not an empty list (matches actual upstream behavior) - Add length validation for the two-list calling convention - Update all tests and docstring examples accordingly Co-Authored-By: Claude Opus 4.6 (1M context) * Consolidate map function tests into parametrized groups Reduce boilerplate by combining make_map construction tests and map accessor function tests into two @pytest.mark.parametrize groups. Co-Authored-By: Claude Opus 4.6 (1M context) * Docstring update Co-authored-by: Nuno Faria * Docstring update Co-authored-by: Nuno Faria * Simplify test for readability Co-authored-by: Nuno Faria * Simplify test for readability Co-authored-by: Nuno Faria --------- Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Nuno Faria --- crates/core/src/functions.rs | 20 +++++ python/datafusion/functions.py | 158 +++++++++++++++++++++++++++++++++ python/tests/test_functions.py | 100 +++++++++++++++++++++ 3 files changed, 278 insertions(+) diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 3f07da95b..5e61b71be 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -93,6 +93,13 @@ fn array_cat(exprs: Vec) -> PyExpr { array_concat(exprs) } +#[pyfunction] +fn make_map(keys: Vec, values: Vec) -> PyExpr { + let keys = keys.into_iter().map(|x| x.into()).collect(); + let values = values.into_iter().map(|x| x.into()).collect(); + datafusion::functions_nested::map::map(keys, values).into() +} + #[pyfunction] #[pyo3(signature = (array, element, index=None))] fn array_position(array: PyExpr, element: PyExpr, index: Option) -> PyExpr { @@ -678,6 +685,12 @@ array_fn!(cardinality, array); array_fn!(flatten, array); array_fn!(range, start stop step); +// Map Functions +array_fn!(map_keys, map); +array_fn!(map_values, map); +array_fn!(map_extract, map key); +array_fn!(map_entries, map); + aggregate_function!(array_agg); aggregate_function!(max); aggregate_function!(min); @@ -1142,6 +1155,13 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(flatten))?; m.add_wrapped(wrap_pyfunction!(cardinality))?; + // Map Functions + m.add_wrapped(wrap_pyfunction!(make_map))?; + m.add_wrapped(wrap_pyfunction!(map_keys))?; + m.add_wrapped(wrap_pyfunction!(map_values))?; + m.add_wrapped(wrap_pyfunction!(map_extract))?; + m.add_wrapped(wrap_pyfunction!(map_entries))?; + // Window Functions m.add_wrapped(wrap_pyfunction!(lead))?; m.add_wrapped(wrap_pyfunction!(lag))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index f1ea3d256..3febb44e3 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -140,6 +140,7 @@ "degrees", "dense_rank", "digest", + "element_at", "empty", "encode", "ends_with", @@ -206,7 +207,12 @@ "make_array", "make_date", "make_list", + "make_map", "make_time", + "map_entries", + "map_extract", + "map_keys", + "map_values", "max", "md5", "mean", @@ -3458,6 +3464,158 @@ def empty(array: Expr) -> Expr: return array_empty(array) +# map functions + + +def make_map(*args: Any) -> Expr: + """Returns a map expression. + + Supports three calling conventions: + + - ``make_map({"a": 1, "b": 2})`` — from a Python dictionary. + - ``make_map([keys], [values])`` — from a list of keys and a list of + their associated values. Both lists must be the same length. + - ``make_map(k1, v1, k2, v2, ...)`` — from alternating keys and their + associated values. + + Keys and values that are not already :py:class:`~datafusion.expr.Expr` + are automatically converted to literal expressions. + + Examples: + From a dictionary: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.make_map({"a": 1, "b": 2}).alias("m")) + >>> result.collect_column("m")[0].as_py() + [('a', 1), ('b', 2)] + + From two lists: + + >>> df = ctx.from_pydict({"key": ["x", "y"], "val": [10, 20]}) + >>> df = df.select( + ... dfn.functions.make_map( + ... [dfn.col("key")], [dfn.col("val")] + ... ).alias("m")) + >>> df.collect_column("m")[0].as_py() + [('x', 10)] + + From alternating keys and values: + + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.make_map("x", 1, "y", 2).alias("m")) + >>> result.collect_column("m")[0].as_py() + [('x', 1), ('y', 2)] + """ + if len(args) == 1 and isinstance(args[0], dict): + key_list = list(args[0].keys()) + value_list = list(args[0].values()) + elif ( + len(args) == 2 # noqa: PLR2004 + and isinstance(args[0], list) + and isinstance(args[1], list) + ): + if len(args[0]) != len(args[1]): + msg = "make_map requires key and value lists to be the same length" + raise ValueError(msg) + key_list = args[0] + value_list = args[1] + elif len(args) >= 2 and len(args) % 2 == 0: # noqa: PLR2004 + key_list = list(args[0::2]) + value_list = list(args[1::2]) + else: + msg = ( + "make_map expects a dict, two lists, or an even number of " + "key-value arguments" + ) + raise ValueError(msg) + + key_exprs = [k if isinstance(k, Expr) else Expr.literal(k) for k in key_list] + val_exprs = [v if isinstance(v, Expr) else Expr.literal(v) for v in value_list] + return Expr(f.make_map([k.expr for k in key_exprs], [v.expr for v in val_exprs])) + + +def map_keys(map: Expr) -> Expr: + """Returns a list of all keys in the map. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> df = df.select( + ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m")) + >>> result = df.select( + ... dfn.functions.map_keys(dfn.col("m")).alias("keys")) + >>> result.collect_column("keys")[0].as_py() + ['x', 'y'] + """ + return Expr(f.map_keys(map.expr)) + + +def map_values(map: Expr) -> Expr: + """Returns a list of all values in the map. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> df = df.select( + ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m")) + >>> result = df.select( + ... dfn.functions.map_values(dfn.col("m")).alias("vals")) + >>> result.collect_column("vals")[0].as_py() + [1, 2] + """ + return Expr(f.map_values(map.expr)) + + +def map_extract(map: Expr, key: Expr) -> Expr: + """Returns the value for a given key in the map. + + Returns ``[None]`` if the key is absent. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> df = df.select( + ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m")) + >>> result = df.select( + ... dfn.functions.map_extract( + ... dfn.col("m"), dfn.lit("x") + ... ).alias("val")) + >>> result.collect_column("val")[0].as_py() + [1] + """ + return Expr(f.map_extract(map.expr, key.expr)) + + +def map_entries(map: Expr) -> Expr: + """Returns a list of all entries (key-value struct pairs) in the map. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> df = df.select( + ... dfn.functions.make_map({"x": 1, "y": 2}).alias("m")) + >>> result = df.select( + ... dfn.functions.map_entries(dfn.col("m")).alias("entries")) + >>> result.collect_column("entries")[0].as_py() + [{'key': 'x', 'value': 1}, {'key': 'y', 'value': 2}] + """ + return Expr(f.map_entries(map.expr)) + + +def element_at(map: Expr, key: Expr) -> Expr: + """Returns the value for a given key in the map. + + Returns ``[None]`` if the key is absent. + + See Also: + This is an alias for :py:func:`map_extract`. + """ + return map_extract(map, key) + + # aggregate functions def approx_distinct( expression: Expr, diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 74fcbffb4..f25c6e78c 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -668,6 +668,106 @@ def test_array_function_obj_tests(stmt, py_expr): assert a == b +@pytest.mark.parametrize( + ("args", "expected"), + [ + pytest.param( + ({"x": 1, "y": 2},), + [("x", 1), ("y", 2)], + id="dict", + ), + pytest.param( + ({"x": literal(1), "y": literal(2)},), + [("x", 1), ("y", 2)], + id="dict_with_exprs", + ), + pytest.param( + ("x", 1, "y", 2), + [("x", 1), ("y", 2)], + id="variadic_pairs", + ), + pytest.param( + (literal("x"), literal(1), literal("y"), literal(2)), + [("x", 1), ("y", 2)], + id="variadic_with_exprs", + ), + ], +) +def test_make_map(args, expected): + ctx = SessionContext() + batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"]) + df = ctx.create_dataframe([[batch]]) + + result = df.select(f.make_map(*args).alias("m")).collect()[0].column(0) + assert result[0].as_py() == expected + + +def test_make_map_from_two_lists(): + ctx = SessionContext() + batch = pa.RecordBatch.from_arrays( + [ + pa.array(["k1", "k2", "k3"]), + pa.array([10, 20, 30]), + ], + names=["keys", "vals"], + ) + df = ctx.create_dataframe([[batch]]) + + m = f.make_map([column("keys")], [column("vals")]) + result = df.select(f.map_keys(m).alias("k")).collect()[0].column(0) + assert result.to_pylist() == [["k1"], ["k2"], ["k3"]] + + result = df.select(f.map_values(m).alias("v")).collect()[0].column(0) + assert result.to_pylist() == [[10], [20], [30]] + + +def test_make_map_odd_args_raises(): + with pytest.raises(ValueError, match="make_map expects"): + f.make_map("x", 1, "y") + + +def test_make_map_mismatched_lengths(): + with pytest.raises(ValueError, match="same length"): + f.make_map(["a", "b"], [1]) + + +@pytest.mark.parametrize( + ("func", "expected"), + [ + pytest.param(f.map_keys, ["x", "y"], id="map_keys"), + pytest.param(f.map_values, [1, 2], id="map_values"), + pytest.param( + lambda m: f.map_extract(m, literal("x")), + [1], + id="map_extract", + ), + pytest.param( + lambda m: f.map_extract(m, literal("z")), + [None], + id="map_extract_missing_key", + ), + pytest.param( + f.map_entries, + [{"key": "x", "value": 1}, {"key": "y", "value": 2}], + id="map_entries", + ), + pytest.param( + lambda m: f.element_at(m, literal("y")), + [2], + id="element_at", + ), + ], +) +def test_map_functions(func, expected): + ctx = SessionContext() + batch = pa.RecordBatch.from_arrays([pa.array([1])], names=["a"]) + df = ctx.create_dataframe([[batch]]) + + m = f.make_map({"x": 1, "y": 2}) + result = df.select(func(m).alias("out")).collect()[0].column(0) + assert result[0].as_py() == expected + + @pytest.mark.parametrize( ("function", "expected_result"), [ From ff15648c5dca6b41d3f6146c6c36c97e605f8561 Mon Sep 17 00:00:00 2001 From: Nuno Faria Date: Sun, 5 Apr 2026 13:29:32 +0100 Subject: [PATCH 09/83] minor: Fix pytest instructions in README (#1477) --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 7c1c71281..7849e7a02 100644 --- a/README.md +++ b/README.md @@ -275,7 +275,7 @@ needing to activate the virtual environment: ```bash uv run --no-project maturin develop --uv -uv run --no-project pytest . +uv run --no-project pytest ``` To run the FFI tests within the examples folder, after you have built From 99bc9602dd077c924685f1fc6e54e6feb3429302 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 6 Apr 2026 07:47:13 -0400 Subject: [PATCH 10/83] Add missing array functions (#1468) * Add missing array/list functions and aliases (#1452) Add new array functions from upstream DataFusion v53: array_any_value, array_distance, array_max, array_min, array_reverse, arrays_zip, string_to_array, and gen_series. Add corresponding list_* aliases and missing list_* aliases for existing functions (list_empty, list_pop_back, list_pop_front, list_has, list_has_all, list_has_any). Also add array_contains/list_contains as aliases for array_has, generate_series as alias for gen_series, and string_to_list as alias for string_to_array. Co-Authored-By: Claude Opus 4.6 (1M context) * Add unit tests for new array/list functions and aliases Tests cover all functions and aliases added in the previous commit: array_any_value, array_distance, array_max, array_min, array_reverse, arrays_zip, string_to_array, gen_series, generate_series, array_contains, list_contains, list_empty, list_pop_back, list_pop_front, list_has, list_has_all, list_has_any, and list_* aliases for the new functions. Co-Authored-By: Claude Opus 4.6 (1M context) * Improve array function APIs: optional params, better naming, restore comment - Make null_string optional in string_to_array/string_to_list - Make step optional in gen_series/generate_series - Rename second_array to element in array_contains/list_has/list_contains - Restore # Window Functions section comment in __all__ - Add tests for optional parameter variants Co-Authored-By: Claude Opus 4.6 (1M context) * Consolidate array/list function tests using pytest parametrize Reduce 26 individual tests to 14 test functions with parametrized cases, eliminating boilerplate while maintaining full coverage. Co-Authored-By: Claude Opus 4.6 (1M context) * Move list alias tests into existing test_array_functions parametrize block Merge standalone tests for list_empty, list_pop_back, list_pop_front, list_has, array_contains, list_contains, list_has_all, and list_has_any into the existing parametrized test_array_functions block alongside their array_* counterparts. Co-Authored-By: Claude Opus 4.6 (1M context) * Merge test_array_any_value into parametrized test_any_value_aliases Use the richer multi-row dataset (including all-nulls case) for both array_any_value and list_any_value via the parametrized test. Co-Authored-By: Claude Opus 4.6 (1M context) * Add arrays_overlap and list_overlap as aliases for array_has_any These aliases match the upstream DataFusion SQL-level aliases, completing the set of missing array functions from issue #1452. Co-Authored-By: Claude Opus 4.6 (1M context) * Add docstring examples for optional params in string_to_array and gen_series Co-Authored-By: Claude Opus 4.6 (1M context) * Update AGENTS file to demonstrate preferred method of documenting python functions --------- Co-authored-by: Claude Opus 4.6 (1M context) --- AGENTS.md | 17 ++ crates/core/src/functions.rs | 56 ++++++ python/datafusion/functions.py | 337 +++++++++++++++++++++++++++++++++ python/tests/test_functions.py | 137 ++++++++++++++ 4 files changed, 547 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index 1853a84cd..f6fdfbd90 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,3 +25,20 @@ 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. + +## Python Function Docstrings + +Every Python function must include a docstring with usage examples. + +- **Examples are required**: Each function needs at least one doctest-style example + demonstrating basic usage. +- **Optional parameters**: If a function has optional parameters, include separate + examples that show usage both without and with the optional arguments. Pass + optional arguments using their keyword name (e.g., `step=dfn.lit(3)`) so readers + can immediately see which parameter is being demonstrated. +- **Reuse input data**: Use the same input data across examples wherever possible. + The examples should demonstrate how different optional arguments change the output + for the same input, making the effect of each option easy to understand. +- **Alias functions**: Functions that are simple aliases (e.g., `list_sort` aliasing + `array_sort`) only need a one-line description and a `See Also` reference to the + primary function. They do not need their own examples. diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 5e61b71be..8bb927718 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -93,6 +93,50 @@ fn array_cat(exprs: Vec) -> PyExpr { array_concat(exprs) } +#[pyfunction] +fn array_distance(array1: PyExpr, array2: PyExpr) -> PyExpr { + let args = vec![array1.into(), array2.into()]; + Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf( + datafusion::functions_nested::distance::array_distance_udf(), + args, + )) + .into() +} + +#[pyfunction] +fn arrays_zip(exprs: Vec) -> PyExpr { + let exprs = exprs.into_iter().map(|x| x.into()).collect(); + datafusion::functions_nested::expr_fn::arrays_zip(exprs).into() +} + +#[pyfunction] +#[pyo3(signature = (string, delimiter, null_string=None))] +fn string_to_array(string: PyExpr, delimiter: PyExpr, null_string: Option) -> PyExpr { + let mut args = vec![string.into(), delimiter.into()]; + if let Some(null_string) = null_string { + args.push(null_string.into()); + } + Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf( + datafusion::functions_nested::string::string_to_array_udf(), + args, + )) + .into() +} + +#[pyfunction] +#[pyo3(signature = (start, stop, step=None))] +fn gen_series(start: PyExpr, stop: PyExpr, step: Option) -> PyExpr { + let mut args = vec![start.into(), stop.into()]; + if let Some(step) = step { + args.push(step.into()); + } + Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf( + datafusion::functions_nested::range::gen_series_udf(), + args, + )) + .into() +} + #[pyfunction] fn make_map(keys: Vec, values: Vec) -> PyExpr { let keys = keys.into_iter().map(|x| x.into()).collect(); @@ -681,6 +725,10 @@ array_fn!(array_intersect, first_array second_array); array_fn!(array_union, array1 array2); array_fn!(array_except, first_array second_array); array_fn!(array_resize, array size value); +array_fn!(array_any_value, array); +array_fn!(array_max, array); +array_fn!(array_min, array); +array_fn!(array_reverse, array); array_fn!(cardinality, array); array_fn!(flatten, array); array_fn!(range, start stop step); @@ -1152,6 +1200,14 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(array_replace_all))?; m.add_wrapped(wrap_pyfunction!(array_sort))?; m.add_wrapped(wrap_pyfunction!(array_slice))?; + m.add_wrapped(wrap_pyfunction!(array_any_value))?; + m.add_wrapped(wrap_pyfunction!(array_distance))?; + m.add_wrapped(wrap_pyfunction!(array_max))?; + m.add_wrapped(wrap_pyfunction!(array_min))?; + m.add_wrapped(wrap_pyfunction!(array_reverse))?; + m.add_wrapped(wrap_pyfunction!(arrays_zip))?; + m.add_wrapped(wrap_pyfunction!(string_to_array))?; + m.add_wrapped(wrap_pyfunction!(gen_series))?; m.add_wrapped(wrap_pyfunction!(flatten))?; m.add_wrapped(wrap_pyfunction!(cardinality))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 3febb44e3..1b267731e 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -53,10 +53,13 @@ "approx_percentile_cont_with_weight", "array", "array_agg", + "array_any_value", "array_append", "array_cat", "array_concat", + "array_contains", "array_dims", + "array_distance", "array_distinct", "array_element", "array_empty", @@ -69,6 +72,8 @@ "array_intersect", "array_join", "array_length", + "array_max", + "array_min", "array_ndims", "array_pop_back", "array_pop_front", @@ -85,10 +90,13 @@ "array_replace_all", "array_replace_n", "array_resize", + "array_reverse", "array_slice", "array_sort", "array_to_string", "array_union", + "arrays_overlap", + "arrays_zip", "arrow_cast", "arrow_typeof", "ascii", @@ -153,6 +161,8 @@ "floor", "from_unixtime", "gcd", + "gen_series", + "generate_series", "greatest", "ifnull", "in_list", @@ -167,19 +177,31 @@ "left", "length", "levenshtein", + "list_any_value", "list_append", "list_cat", "list_concat", + "list_contains", "list_dims", + "list_distance", "list_distinct", "list_element", + "list_empty", "list_except", "list_extract", + "list_has", + "list_has_all", + "list_has_any", "list_indexof", "list_intersect", "list_join", "list_length", + "list_max", + "list_min", "list_ndims", + "list_overlap", + "list_pop_back", + "list_pop_front", "list_position", "list_positions", "list_prepend", @@ -193,10 +215,12 @@ "list_replace_all", "list_replace_n", "list_resize", + "list_reverse", "list_slice", "list_sort", "list_to_string", "list_union", + "list_zip", "ln", "log", "log2", @@ -273,6 +297,8 @@ "stddev_pop", "stddev_samp", "string_agg", + "string_to_array", + "string_to_list", "strpos", "struct", "substr", @@ -2794,6 +2820,15 @@ def array_empty(array: Expr) -> Expr: return Expr(f.array_empty(array.expr)) +def list_empty(array: Expr) -> Expr: + """Returns a boolean indicating whether the array is empty. + + See Also: + This is an alias for :py:func:`array_empty`. + """ + return array_empty(array) + + def array_extract(array: Expr, n: Expr) -> Expr: """Extracts the element with the index n from the array. @@ -2891,6 +2926,69 @@ def array_has_any(first_array: Expr, second_array: Expr) -> Expr: return Expr(f.array_has_any(first_array.expr, second_array.expr)) +def array_contains(array: Expr, element: Expr) -> Expr: + """Returns true if the element appears in the array, otherwise false. + + See Also: + This is an alias for :py:func:`array_has`. + """ + return array_has(array, element) + + +def list_has(array: Expr, element: Expr) -> Expr: + """Returns true if the element appears in the array, otherwise false. + + See Also: + This is an alias for :py:func:`array_has`. + """ + return array_has(array, element) + + +def list_has_all(first_array: Expr, second_array: Expr) -> Expr: + """Determines if there is complete overlap ``second_array`` in ``first_array``. + + See Also: + This is an alias for :py:func:`array_has_all`. + """ + return array_has_all(first_array, second_array) + + +def list_has_any(first_array: Expr, second_array: Expr) -> Expr: + """Determine if there is an overlap between ``first_array`` and ``second_array``. + + See Also: + This is an alias for :py:func:`array_has_any`. + """ + return array_has_any(first_array, second_array) + + +def arrays_overlap(first_array: Expr, second_array: Expr) -> Expr: + """Returns true if any element appears in both arrays. + + See Also: + This is an alias for :py:func:`array_has_any`. + """ + return array_has_any(first_array, second_array) + + +def list_overlap(first_array: Expr, second_array: Expr) -> Expr: + """Returns true if any element appears in both arrays. + + See Also: + This is an alias for :py:func:`array_has_any`. + """ + return array_has_any(first_array, second_array) + + +def list_contains(array: Expr, element: Expr) -> Expr: + """Returns true if the element appears in the array, otherwise false. + + See Also: + This is an alias for :py:func:`array_has`. + """ + return array_has(array, element) + + def array_position(array: Expr, element: Expr, index: int | None = 1) -> Expr: """Return the position of the first occurrence of ``element`` in ``array``. @@ -3058,6 +3156,24 @@ def array_pop_front(array: Expr) -> Expr: return Expr(f.array_pop_front(array.expr)) +def list_pop_back(array: Expr) -> Expr: + """Returns the array without the last element. + + See Also: + This is an alias for :py:func:`array_pop_back`. + """ + return array_pop_back(array) + + +def list_pop_front(array: Expr) -> Expr: + """Returns the array without the first element. + + See Also: + This is an alias for :py:func:`array_pop_front`. + """ + return array_pop_front(array) + + def array_remove(array: Expr, element: Expr) -> Expr: """Removes the first element from the array equal to the given value. @@ -3429,6 +3545,227 @@ def list_resize(array: Expr, size: Expr, value: Expr) -> Expr: return array_resize(array, size, value) +def array_any_value(array: Expr) -> Expr: + """Returns the first non-null element in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[None, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_any_value(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 2 + """ + return Expr(f.array_any_value(array.expr)) + + +def list_any_value(array: Expr) -> Expr: + """Returns the first non-null element in the array. + + See Also: + This is an alias for :py:func:`array_any_value`. + """ + return array_any_value(array) + + +def array_distance(array1: Expr, array2: Expr) -> Expr: + """Returns the Euclidean distance between two numeric arrays. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 4.0]]}) + >>> result = df.select( + ... dfn.functions.array_distance( + ... dfn.col("a"), dfn.col("b"), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + 2.0 + """ + return Expr(f.array_distance(array1.expr, array2.expr)) + + +def list_distance(array1: Expr, array2: Expr) -> Expr: + """Returns the Euclidean distance between two numeric arrays. + + See Also: + This is an alias for :py:func:`array_distance`. + """ + return array_distance(array1, array2) + + +def array_max(array: Expr) -> Expr: + """Returns the maximum value in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_max(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 3 + """ + return Expr(f.array_max(array.expr)) + + +def list_max(array: Expr) -> Expr: + """Returns the maximum value in the array. + + See Also: + This is an alias for :py:func:`array_max`. + """ + return array_max(array) + + +def array_min(array: Expr) -> Expr: + """Returns the minimum value in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_min(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 1 + """ + return Expr(f.array_min(array.expr)) + + +def list_min(array: Expr) -> Expr: + """Returns the minimum value in the array. + + See Also: + This is an alias for :py:func:`array_min`. + """ + return array_min(array) + + +def array_reverse(array: Expr) -> Expr: + """Reverses the order of elements in the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_reverse(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [3, 2, 1] + """ + return Expr(f.array_reverse(array.expr)) + + +def list_reverse(array: Expr) -> Expr: + """Reverses the order of elements in the array. + + See Also: + This is an alias for :py:func:`array_reverse`. + """ + return array_reverse(array) + + +def arrays_zip(*arrays: Expr) -> Expr: + """Combines multiple arrays into a single array of structs. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]}) + >>> result = df.select( + ... dfn.functions.arrays_zip(dfn.col("a"), dfn.col("b")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [{'c0': 1, 'c1': 3}, {'c0': 2, 'c1': 4}] + """ + args = [a.expr for a in arrays] + return Expr(f.arrays_zip(args)) + + +def list_zip(*arrays: Expr) -> Expr: + """Combines multiple arrays into a single array of structs. + + See Also: + This is an alias for :py:func:`arrays_zip`. + """ + return arrays_zip(*arrays) + + +def string_to_array( + string: Expr, delimiter: Expr, null_string: Expr | None = None +) -> Expr: + """Splits a string based on a delimiter and returns an array of parts. + + Any parts matching the optional ``null_string`` will be replaced with ``NULL``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello,world"]}) + >>> result = df.select( + ... dfn.functions.string_to_array( + ... dfn.col("a"), dfn.lit(","), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + ['hello', 'world'] + + Replace parts matching a ``null_string`` with ``NULL``: + + >>> result = df.select( + ... dfn.functions.string_to_array( + ... dfn.col("a"), dfn.lit(","), null_string=dfn.lit("world"), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + ['hello', None] + """ + null_expr = null_string.expr if null_string is not None else None + return Expr(f.string_to_array(string.expr, delimiter.expr, null_expr)) + + +def string_to_list( + string: Expr, delimiter: Expr, null_string: Expr | None = None +) -> Expr: + """Splits a string based on a delimiter and returns an array of parts. + + See Also: + This is an alias for :py:func:`string_to_array`. + """ + return string_to_array(string, delimiter, null_string) + + +def gen_series(start: Expr, stop: Expr, step: Expr | None = None) -> Expr: + """Creates a list of values in the range between start and stop. + + Unlike :py:func:`range`, this includes the upper bound. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0]}) + >>> result = df.select( + ... dfn.functions.gen_series( + ... dfn.lit(1), dfn.lit(5), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2, 3, 4, 5] + + Specify a custom ``step``: + + >>> result = df.select( + ... dfn.functions.gen_series( + ... dfn.lit(1), dfn.lit(10), step=dfn.lit(3), + ... ).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 4, 7, 10] + """ + step_expr = step.expr if step is not None else None + return Expr(f.gen_series(start.expr, stop.expr, step_expr)) + + +def generate_series(start: Expr, stop: Expr, step: Expr | None = None) -> Expr: + """Creates a list of values in the range between start and stop. + + Unlike :py:func:`range`, this includes the upper bound. + + See Also: + This is an alias for :py:func:`gen_series`. + """ + return gen_series(start, stop, step) + + def flatten(array: Expr) -> Expr: """Flattens an array of arrays into a single array. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index f25c6e78c..2100da9ae 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -330,6 +330,10 @@ def py_flatten(arr): f.empty, lambda data: [len(r) == 0 for r in data], ), + ( + f.list_empty, + lambda data: [len(r) == 0 for r in data], + ), ( lambda col: f.array_extract(col, literal(1)), lambda data: [r[0] for r in data], @@ -354,18 +358,54 @@ def py_flatten(arr): lambda col: f.array_has(col, literal(1.0)), lambda data: [1.0 in r for r in data], ), + ( + lambda col: f.list_has(col, literal(1.0)), + lambda data: [1.0 in r for r in data], + ), + ( + lambda col: f.array_contains(col, literal(1.0)), + lambda data: [1.0 in r for r in data], + ), + ( + lambda col: f.list_contains(col, literal(1.0)), + lambda data: [1.0 in r for r in data], + ), ( lambda col: f.array_has_all( col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]]) ), lambda data: [np.all([v in r for v in [1.0, 3.0, 5.0]]) for r in data], ), + ( + lambda col: f.list_has_all( + col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]]) + ), + lambda data: [np.all([v in r for v in [1.0, 3.0, 5.0]]) for r in data], + ), ( lambda col: f.array_has_any( col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]]) ), lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data], ), + ( + lambda col: f.list_has_any( + col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]]) + ), + lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data], + ), + ( + lambda col: f.arrays_overlap( + col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]]) + ), + lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data], + ), + ( + lambda col: f.list_overlap( + col, f.make_array(*[literal(v) for v in [1.0, 3.0, 5.0]]) + ), + lambda data: [np.any([v in r for v in [1.0, 3.0, 5.0]]) for r in data], + ), ( lambda col: f.array_position(col, literal(1.0)), lambda data: [py_indexof(r, 1.0) for r in data], @@ -418,10 +458,18 @@ def py_flatten(arr): f.array_pop_back, lambda data: [arr[:-1] for arr in data], ), + ( + f.list_pop_back, + lambda data: [arr[:-1] for arr in data], + ), ( f.array_pop_front, lambda data: [arr[1:] for arr in data], ), + ( + f.list_pop_front, + lambda data: [arr[1:] for arr in data], + ), ( lambda col: f.array_remove(col, literal(3.0)), lambda data: [py_arr_remove(arr, 3.0, 1) for arr in data], @@ -1760,3 +1808,92 @@ def df_with_nulls(): def test_conditional_functions(df_with_nulls, expr, expected): result = df_with_nulls.select(expr.alias("result")).collect()[0] assert result.column(0) == expected + + +@pytest.mark.parametrize("func", [f.array_any_value, f.list_any_value]) +def test_any_value_aliases(func): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[None, 2, 3], [None, None, None], [1, 2, 3]]}) + result = df.select(func(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values[0] == 2 + assert values[1] is None + assert values[2] == 1 + + +@pytest.mark.parametrize("func", [f.array_distance, f.list_distance]) +def test_array_distance_aliases(func): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 4.0]]}) + result = df.select(func(column("a"), column("b")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == pytest.approx(2.0) + + +@pytest.mark.parametrize( + ("func", "expected"), + [ + (f.array_max, [5, 10]), + (f.list_max, [5, 10]), + (f.array_min, [1, 2]), + (f.list_min, [1, 2]), + ], +) +def test_array_min_max(func, expected): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 5, 3], [10, 2]]}) + result = df.select(func(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values == expected + + +@pytest.mark.parametrize("func", [f.array_reverse, f.list_reverse]) +def test_array_reverse_aliases(func): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3], [4, 5]]}) + result = df.select(func(column("a")).alias("v")).collect() + values = [row.as_py() for row in result[0].column(0)] + assert values == [[3, 2, 1], [5, 4]] + + +@pytest.mark.parametrize("func", [f.arrays_zip, f.list_zip]) +def test_arrays_zip_aliases(func): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]}) + result = df.select(func(column("a"), column("b")).alias("v")).collect() + values = result[0].column(0)[0].as_py() + assert values == [{"c0": 1, "c1": 3}, {"c0": 2, "c1": 4}] + + +@pytest.mark.parametrize("func", [f.string_to_array, f.string_to_list]) +def test_string_to_array_aliases(func): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["hello,world,foo"]}) + result = df.select(func(column("a"), literal(",")).alias("v")).collect() + assert result[0].column(0)[0].as_py() == ["hello", "world", "foo"] + + +def test_string_to_array_with_null_string(): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["hello,NA,world"]}) + result = df.select( + f.string_to_array(column("a"), literal(","), literal("NA")).alias("v") + ).collect() + values = result[0].column(0)[0].as_py() + assert values == ["hello", None, "world"] + + +@pytest.mark.parametrize("func", [f.gen_series, f.generate_series]) +def test_gen_series_aliases(func): + ctx = SessionContext() + df = ctx.from_pydict({"a": [0]}) + result = df.select(func(literal(1), literal(5)).alias("v")).collect() + assert result[0].column(0)[0].as_py() == [1, 2, 3, 4, 5] + + +def test_gen_series_with_step(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [0]}) + result = df.select( + f.gen_series(literal(1), literal(10), literal(3)).alias("v") + ).collect() + assert result[0].column(0)[0].as_py() == [1, 4, 7, 10] From d07fdb3ef7d211920f40d0106fa50161c0bf20ce Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 6 Apr 2026 08:54:30 -0400 Subject: [PATCH 11/83] Add missing scalar functions (#1470) * Add missing scalar functions: get_field, union_extract, union_tag, arrow_metadata, version, row Expose upstream DataFusion scalar functions that were not yet available in the Python API. Closes #1453. - get_field: extracts a field from a struct or map by name - union_extract: extracts a value from a union type by field name - union_tag: returns the active field name of a union type - arrow_metadata: returns Arrow field metadata (all or by key) - version: returns the DataFusion version string - row: alias for the struct constructor Note: arrow_try_cast was listed in the issue but does not exist in DataFusion 53, so it is not included. Co-Authored-By: Claude Opus 4.6 (1M context) * Add tests for new scalar functions Tests for get_field, arrow_metadata, version, row, union_tag, and union_extract. Co-Authored-By: Claude Opus 4.6 (1M context) * Accept str for field name and type parameters in scalar functions Allow arrow_cast, get_field, and union_extract to accept plain str arguments instead of requiring Expr wrappers. Also improve arrow_metadata test coverage and fix parameter shadowing. Co-Authored-By: Claude Opus 4.6 (1M context) * Accept str for key parameter in arrow_metadata for consistency Co-Authored-By: Claude Opus 4.6 (1M context) * Add doctest examples and fix docstring style for new scalar functions Replace Args/Returns sections with doctest Examples blocks for arrow_metadata, get_field, union_extract, union_tag, and version to match existing codebase conventions. Simplify row to alias-style docstring with See Also reference. Document that arrow_cast accepts both str and Expr for data_type. Co-Authored-By: Claude Opus 4.6 (1M context) * Support pyarrow DataType in arrow_cast Allow arrow_cast to accept a pyarrow DataType in addition to str and Expr. The DataType is converted to its string representation before being passed to DataFusion. Adds test coverage for the new input type. Co-Authored-By: Claude Opus 4.6 (1M context) * Document bracket syntax shorthand in get_field docstring Note that expr["field"] is a convenient alternative when the field name is a static string, and get_field is needed for dynamic expressions. Add a second doctest example showing the bracket syntax. Co-Authored-By: Claude Opus 4.6 (1M context) * Fix arrow_cast with pyarrow DataType by delegating to Expr.cast Use the existing Rust-side PyArrowType conversion via Expr.cast() instead of str() which produces pyarrow type names that DataFusion does not recognize. Co-Authored-By: Claude Opus 4.6 (1M context) * Clarify when to use arrow_cast vs Expr.cast in docstring Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- crates/core/src/functions.rs | 26 +++++ python/datafusion/functions.py | 174 ++++++++++++++++++++++++++++++++- python/tests/test_functions.py | 105 ++++++++++++++++++-- 3 files changed, 296 insertions(+), 9 deletions(-) diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 8bb927718..74654ce46 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -695,8 +695,29 @@ expr_fn_vec!(named_struct); expr_fn!(from_unixtime, unixtime); expr_fn!(arrow_typeof, arg_1); expr_fn!(arrow_cast, arg_1 datatype); +expr_fn_vec!(arrow_metadata); +expr_fn!(union_tag, arg1); expr_fn!(random); +#[pyfunction] +fn get_field(expr: PyExpr, name: PyExpr) -> PyExpr { + functions::core::get_field() + .call(vec![expr.into(), name.into()]) + .into() +} + +#[pyfunction] +fn union_extract(union_expr: PyExpr, field_name: PyExpr) -> PyExpr { + functions::core::union_extract() + .call(vec![union_expr.into(), field_name.into()]) + .into() +} + +#[pyfunction] +fn version() -> PyExpr { + functions::core::version().call(vec![]).into() +} + // Array Functions array_fn!(array_append, array element); array_fn!(array_to_string, array delimiter); @@ -1014,6 +1035,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(array_agg))?; m.add_wrapped(wrap_pyfunction!(arrow_typeof))?; m.add_wrapped(wrap_pyfunction!(arrow_cast))?; + m.add_wrapped(wrap_pyfunction!(arrow_metadata))?; m.add_wrapped(wrap_pyfunction!(ascii))?; m.add_wrapped(wrap_pyfunction!(asin))?; m.add_wrapped(wrap_pyfunction!(asinh))?; @@ -1142,6 +1164,10 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(trim))?; m.add_wrapped(wrap_pyfunction!(trunc))?; m.add_wrapped(wrap_pyfunction!(upper))?; + m.add_wrapped(wrap_pyfunction!(get_field))?; + m.add_wrapped(wrap_pyfunction!(union_extract))?; + m.add_wrapped(wrap_pyfunction!(union_tag))?; + m.add_wrapped(wrap_pyfunction!(version))?; m.add_wrapped(wrap_pyfunction!(self::uuid))?; // Use self to avoid name collision m.add_wrapped(wrap_pyfunction!(var_pop))?; m.add_wrapped(wrap_pyfunction!(var_sample))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 1b267731e..aa7f28746 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -98,6 +98,7 @@ "arrays_overlap", "arrays_zip", "arrow_cast", + "arrow_metadata", "arrow_typeof", "ascii", "asin", @@ -163,6 +164,7 @@ "gcd", "gen_series", "generate_series", + "get_field", "greatest", "ifnull", "in_list", @@ -280,6 +282,7 @@ "reverse", "right", "round", + "row", "row_number", "rpad", "rtrim", @@ -322,12 +325,15 @@ "translate", "trim", "trunc", + "union_extract", + "union_tag", "upper", "uuid", "var", "var_pop", "var_samp", "var_sample", + "version", "when", # Window Functions "window", @@ -2628,22 +2634,184 @@ def arrow_typeof(arg: Expr) -> Expr: return Expr(f.arrow_typeof(arg.expr)) -def arrow_cast(expr: Expr, data_type: Expr) -> Expr: +def arrow_cast(expr: Expr, data_type: Expr | str | pa.DataType) -> Expr: """Casts an expression to a specified data type. + The ``data_type`` can be a string, a ``pyarrow.DataType``, or an + ``Expr``. For simple types, :py:meth:`Expr.cast() + ` is more concise + (e.g., ``col("a").cast(pa.float64())``). Use ``arrow_cast`` when + you want to specify the target type as a string using DataFusion's + type syntax, which can be more readable for complex types like + ``"Timestamp(Nanosecond, None)"``. + Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [1]}) - >>> data_type = dfn.string_literal("Float64") >>> result = df.select( - ... dfn.functions.arrow_cast(dfn.col("a"), data_type).alias("c") + ... dfn.functions.arrow_cast(dfn.col("a"), "Float64").alias("c") + ... ) + >>> result.collect_column("c")[0].as_py() + 1.0 + + >>> import pyarrow as pa + >>> result = df.select( + ... dfn.functions.arrow_cast( + ... dfn.col("a"), data_type=pa.float64() + ... ).alias("c") ... ) >>> result.collect_column("c")[0].as_py() 1.0 """ + if isinstance(data_type, pa.DataType): + return expr.cast(data_type) + if isinstance(data_type, str): + data_type = Expr.string_literal(data_type) return Expr(f.arrow_cast(expr.expr, data_type.expr)) +def arrow_metadata(expr: Expr, key: Expr | str | None = None) -> Expr: + """Returns the metadata of the input expression. + + If called with one argument, returns a Map of all metadata key-value pairs. + If called with two arguments, returns the value for the specified metadata key. + + Examples: + >>> import pyarrow as pa + >>> field = pa.field("val", pa.int64(), metadata={"k": "v"}) + >>> schema = pa.schema([field]) + >>> batch = pa.RecordBatch.from_arrays([pa.array([1])], schema=schema) + >>> ctx = dfn.SessionContext() + >>> df = ctx.create_dataframe([[batch]]) + >>> result = df.select( + ... dfn.functions.arrow_metadata(dfn.col("val")).alias("meta") + ... ) + >>> ("k", "v") in result.collect_column("meta")[0].as_py() + True + + >>> result = df.select( + ... dfn.functions.arrow_metadata( + ... dfn.col("val"), key="k" + ... ).alias("meta_val") + ... ) + >>> result.collect_column("meta_val")[0].as_py() + 'v' + """ + if key is None: + return Expr(f.arrow_metadata(expr.expr)) + if isinstance(key, str): + key = Expr.string_literal(key) + return Expr(f.arrow_metadata(expr.expr, key.expr)) + + +def get_field(expr: Expr, name: Expr | str) -> Expr: + """Extracts a field from a struct or map by name. + + When the field name is a static string, the bracket operator + ``expr["field"]`` is a convenient shorthand. Use ``get_field`` + when the field name is a dynamic expression. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1], "b": [2]}) + >>> df = df.with_column( + ... "s", + ... dfn.functions.named_struct( + ... [("x", dfn.col("a")), ("y", dfn.col("b"))] + ... ), + ... ) + >>> result = df.select( + ... dfn.functions.get_field(dfn.col("s"), "x").alias("x_val") + ... ) + >>> result.collect_column("x_val")[0].as_py() + 1 + + Equivalent using bracket syntax: + + >>> result = df.select( + ... dfn.col("s")["x"].alias("x_val") + ... ) + >>> result.collect_column("x_val")[0].as_py() + 1 + """ + if isinstance(name, str): + name = Expr.string_literal(name) + return Expr(f.get_field(expr.expr, name.expr)) + + +def union_extract(union_expr: Expr, field_name: Expr | str) -> Expr: + """Extracts a value from a union type by field name. + + Returns the value of the named field if it is the currently selected + variant, otherwise returns NULL. + + Examples: + >>> import pyarrow as pa + >>> ctx = dfn.SessionContext() + >>> types = pa.array([0, 1, 0], type=pa.int8()) + >>> offsets = pa.array([0, 0, 1], type=pa.int32()) + >>> arr = pa.UnionArray.from_dense( + ... types, offsets, [pa.array([1, 2]), pa.array(["hi"])], + ... ["int", "str"], [0, 1], + ... ) + >>> batch = pa.RecordBatch.from_arrays([arr], names=["u"]) + >>> df = ctx.create_dataframe([[batch]]) + >>> result = df.select( + ... dfn.functions.union_extract(dfn.col("u"), "int").alias("val") + ... ) + >>> result.collect_column("val").to_pylist() + [1, None, 2] + """ + if isinstance(field_name, str): + field_name = Expr.string_literal(field_name) + return Expr(f.union_extract(union_expr.expr, field_name.expr)) + + +def union_tag(union_expr: Expr) -> Expr: + """Returns the tag (active field name) of a union type. + + Examples: + >>> import pyarrow as pa + >>> ctx = dfn.SessionContext() + >>> types = pa.array([0, 1, 0], type=pa.int8()) + >>> offsets = pa.array([0, 0, 1], type=pa.int32()) + >>> arr = pa.UnionArray.from_dense( + ... types, offsets, [pa.array([1, 2]), pa.array(["hi"])], + ... ["int", "str"], [0, 1], + ... ) + >>> batch = pa.RecordBatch.from_arrays([arr], names=["u"]) + >>> df = ctx.create_dataframe([[batch]]) + >>> result = df.select( + ... dfn.functions.union_tag(dfn.col("u")).alias("tag") + ... ) + >>> result.collect_column("tag").to_pylist() + ['int', 'str', 'int'] + """ + return Expr(f.union_tag(union_expr.expr)) + + +def version() -> Expr: + """Returns the DataFusion version string. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.empty_table() + >>> result = df.select(dfn.functions.version().alias("v")) + >>> "Apache DataFusion" in result.collect_column("v")[0].as_py() + True + """ + return Expr(f.version()) + + +def row(*args: Expr) -> Expr: + """Returns a struct with the given arguments. + + See Also: + This is an alias for :py:func:`struct`. + """ + return struct(*args) + + def random() -> Expr: """Returns a random value in the range ``0.0 <= x < 1.0``. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 2100da9ae..4e99fa9e3 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -20,7 +20,7 @@ import numpy as np import pyarrow as pa import pytest -from datafusion import SessionContext, column, literal, string_literal +from datafusion import SessionContext, column, literal from datafusion import functions as f np.seterr(invalid="ignore") @@ -1291,11 +1291,8 @@ def test_make_time(df): def test_arrow_cast(df): df = df.select( - # we use `string_literal` to return utf8 instead of `literal` which returns - # utf8view because datafusion.arrow_cast expects a utf8 instead of utf8view - # https://github.com/apache/datafusion/blob/86740bfd3d9831d6b7c1d0e1bf4a21d91598a0ac/datafusion/functions/src/core/arrow_cast.rs#L179 - f.arrow_cast(column("b"), string_literal("Float64")).alias("b_as_float"), - f.arrow_cast(column("b"), string_literal("Int32")).alias("b_as_int"), + f.arrow_cast(column("b"), "Float64").alias("b_as_float"), + f.arrow_cast(column("b"), "Int32").alias("b_as_int"), ) result = df.collect() assert len(result) == 1 @@ -1305,6 +1302,19 @@ def test_arrow_cast(df): assert result.column(1) == pa.array([4, 5, 6], type=pa.int32()) +def test_arrow_cast_with_pyarrow_type(df): + df = df.select( + f.arrow_cast(column("b"), pa.float64()).alias("b_as_float"), + f.arrow_cast(column("b"), pa.int32()).alias("b_as_int"), + f.arrow_cast(column("b"), pa.string()).alias("b_as_str"), + ) + result = df.collect()[0] + + assert result.column(0) == pa.array([4.0, 5.0, 6.0], type=pa.float64()) + assert result.column(1) == pa.array([4, 5, 6], type=pa.int32()) + assert result.column(2) == pa.array(["4", "5", "6"], type=pa.string()) + + def test_case(df): df = df.select( f.case(column("b")).when(literal(4), literal(10)).otherwise(literal(8)), @@ -1810,6 +1820,89 @@ def test_conditional_functions(df_with_nulls, expr, expected): assert result.column(0) == expected +def test_get_field(df): + df = df.with_column( + "s", + f.named_struct( + [ + ("x", column("a")), + ("y", column("b")), + ] + ), + ) + result = df.select( + f.get_field(column("s"), "x").alias("x_val"), + f.get_field(column("s"), "y").alias("y_val"), + ).collect()[0] + + assert result.column(0) == pa.array(["Hello", "World", "!"], type=pa.string_view()) + assert result.column(1) == pa.array([4, 5, 6]) + + +def test_arrow_metadata(): + ctx = SessionContext() + field = pa.field("val", pa.int64(), metadata={"key1": "value1", "key2": "value2"}) + schema = pa.schema([field]) + batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], schema=schema) + df = ctx.create_dataframe([[batch]]) + + # One-argument form: returns a Map of all metadata key-value pairs + result = df.select( + f.arrow_metadata(column("val")).alias("meta"), + ).collect()[0] + assert result.column(0).type == pa.map_(pa.utf8(), pa.utf8()) + meta = result.column(0)[0].as_py() + assert ("key1", "value1") in meta + assert ("key2", "value2") in meta + + # Two-argument form: returns the value for a specific metadata key + result = df.select( + f.arrow_metadata(column("val"), "key1").alias("meta_val"), + ).collect()[0] + assert result.column(0)[0].as_py() == "value1" + + +def test_version(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1]}) + result = df.select(f.version().alias("v")).collect()[0] + version_str = result.column(0)[0].as_py() + assert "Apache DataFusion" in version_str + + +def test_row(df): + result = df.select( + f.row(column("a"), column("b")).alias("r"), + f.struct(column("a"), column("b")).alias("s"), + ).collect()[0] + # row is an alias for struct, so they should produce the same output + assert result.column(0) == result.column(1) + + +def test_union_tag(): + ctx = SessionContext() + types = pa.array([0, 1, 0], type=pa.int8()) + offsets = pa.array([0, 0, 1], type=pa.int32()) + children = [pa.array([1, 2]), pa.array(["hello"])] + arr = pa.UnionArray.from_dense(types, offsets, children, ["int", "str"], [0, 1]) + df = ctx.create_dataframe([[pa.RecordBatch.from_arrays([arr], names=["u"])]]) + + result = df.select(f.union_tag(column("u")).alias("tag")).collect()[0] + assert result.column(0).to_pylist() == ["int", "str", "int"] + + +def test_union_extract(): + ctx = SessionContext() + types = pa.array([0, 1, 0], type=pa.int8()) + offsets = pa.array([0, 0, 1], type=pa.int32()) + children = [pa.array([1, 2]), pa.array(["hello"])] + arr = pa.UnionArray.from_dense(types, offsets, children, ["int", "str"], [0, 1]) + df = ctx.create_dataframe([[pa.RecordBatch.from_arrays([arr], names=["u"])]]) + + result = df.select(f.union_extract(column("u"), "int").alias("val")).collect()[0] + assert result.column(0).to_pylist() == [1, None, 2] + + @pytest.mark.parametrize("func", [f.array_any_value, f.list_any_value]) def test_any_value_aliases(func): ctx = SessionContext() From 898d73de20346bba7241907bb18cba47da53e9a9 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 7 Apr 2026 09:01:36 -0400 Subject: [PATCH 12/83] Add missing aggregate functions (#1471) * Add missing aggregate functions: grouping, percentile_cont, var_population Expose upstream DataFusion aggregate functions that were not yet available in the Python API. Closes #1454. - grouping: returns grouping set membership indicator (rewritten by the ResolveGroupingFunction analyzer rule before physical planning) - percentile_cont: computes exact percentile using continuous interpolation (unlike approx_percentile_cont which uses t-digest) - var_population: alias for var_pop Co-Authored-By: Claude Opus 4.6 (1M context) * Fix grouping() distinct parameter type for API consistency Co-Authored-By: Claude Opus 4.6 (1M context) * Improve aggregate function tests and docstrings per review feedback Add docstring example to grouping(), parametrize percentile_cont tests, and add multi-column grouping test case. Co-Authored-By: Claude Opus 4.6 (1M context) * Add GroupingSet.rollup, .cube, and .grouping_sets factory methods Expose ROLLUP, CUBE, and GROUPING SETS via the DataFrame API by adding static methods on GroupingSet that construct the corresponding Expr variants. Update grouping() docstring and tests to use the new API. Co-Authored-By: Claude Opus 4.6 (1M context) * Remove _GroupingSetInternal alias, use expr_internal.GroupingSet directly Co-Authored-By: Claude Opus 4.6 (1M context) * Parametrize grouping set tests for rollup and cube Co-Authored-By: Claude Opus 4.6 (1M context) * Add grouping sets documentation and note grouping() alias limitation Add user documentation for GroupingSet.rollup, .cube, and .grouping_sets with Pokemon dataset examples. Document the upstream alias limitation (apache/datafusion#21411) in both the grouping() docstring and the aggregation user guide. Co-Authored-By: Claude Opus 4.6 (1M context) * Add grouping sets note to DataFrame.aggregate() docstring Co-Authored-By: Claude Opus 4.6 (1M context) * Address PR review feedback: add quantile_cont alias and simplify examples - Add quantile_cont as alias for percentile_cont (matches upstream) - Replace pa.concat_arrays batch pattern with collect_column() in docstrings - Add percentile_cont, quantile_cont, var_population to docs function list Co-Authored-By: Claude Opus 4.6 (1M context) * Accept string column names in GroupingSet factory methods GroupingSet.rollup(), .cube(), and .grouping_sets() now accept both Expr objects and string column names, consistent with DataFrame.aggregate(). Co-Authored-By: Claude Opus 4.6 (1M context) * Add agent instructions to keep aggregation/window docs in sync Co-Authored-By: Claude Opus 4.6 (1M context) * dfn is already available globally * Remove unnecessary import on doctest --------- Co-authored-by: Claude Opus 4.6 (1M context) --- AGENTS.md | 12 ++ crates/core/src/expr/grouping_set.rs | 37 +++- crates/core/src/functions.rs | 23 ++- .../common-operations/aggregations.rst | 172 +++++++++++++++++- python/datafusion/dataframe.py | 16 +- python/datafusion/expr.py | 127 ++++++++++++- python/datafusion/functions.py | 130 ++++++++++++- python/tests/test_functions.py | 109 +++++++++++ 8 files changed, 614 insertions(+), 12 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f6fdfbd90..86c2e9c3b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -42,3 +42,15 @@ Every Python function must include a docstring with usage examples. - **Alias functions**: Functions that are simple aliases (e.g., `list_sort` aliasing `array_sort`) only need a one-line description and a `See Also` reference to the primary function. They do not need their own examples. + +## Aggregate and Window Function Documentation + +When adding or updating an aggregate or window function, ensure the corresponding +site documentation is kept in sync: + +- **Aggregations**: `docs/source/user-guide/common-operations/aggregations.rst` — + add new aggregate functions to the "Aggregate Functions" list and include usage + examples if appropriate. +- **Window functions**: `docs/source/user-guide/common-operations/windows.rst` — + add new window functions to the "Available Functions" list and include usage + examples if appropriate. diff --git a/crates/core/src/expr/grouping_set.rs b/crates/core/src/expr/grouping_set.rs index 549a866ed..11d8f4fcd 100644 --- a/crates/core/src/expr/grouping_set.rs +++ b/crates/core/src/expr/grouping_set.rs @@ -15,9 +15,11 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::GroupingSet; +use datafusion::logical_expr::{Expr, GroupingSet}; use pyo3::prelude::*; +use crate::expr::PyExpr; + #[pyclass( from_py_object, frozen, @@ -30,6 +32,39 @@ pub struct PyGroupingSet { grouping_set: GroupingSet, } +#[pymethods] +impl PyGroupingSet { + #[staticmethod] + #[pyo3(signature = (*exprs))] + fn rollup(exprs: Vec) -> PyExpr { + Expr::GroupingSet(GroupingSet::Rollup( + exprs.into_iter().map(|e| e.expr).collect(), + )) + .into() + } + + #[staticmethod] + #[pyo3(signature = (*exprs))] + fn cube(exprs: Vec) -> PyExpr { + Expr::GroupingSet(GroupingSet::Cube( + exprs.into_iter().map(|e| e.expr).collect(), + )) + .into() + } + + #[staticmethod] + #[pyo3(signature = (*expr_lists))] + fn grouping_sets(expr_lists: Vec>) -> PyExpr { + Expr::GroupingSet(GroupingSet::GroupingSets( + expr_lists + .into_iter() + .map(|list| list.into_iter().map(|e| e.expr).collect()) + .collect(), + )) + .into() + } +} + impl From for GroupingSet { fn from(grouping_set: PyGroupingSet) -> Self { grouping_set.grouping_set diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 74654ce46..f173aaa51 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -791,9 +791,10 @@ aggregate_function!(var_pop); aggregate_function!(approx_distinct); aggregate_function!(approx_median); -// Code is commented out since grouping is not yet implemented -// https://github.com/apache/datafusion-python/issues/861 -// aggregate_function!(grouping); +// The grouping function's physical plan is not implemented, but the +// ResolveGroupingFunction analyzer rule rewrites it before the physical +// planner sees it, so it works correctly at runtime. +aggregate_function!(grouping); #[pyfunction] #[pyo3(signature = (sort_expression, percentile, num_centroids=None, filter=None))] @@ -831,6 +832,19 @@ pub fn approx_percentile_cont_with_weight( add_builder_fns_to_aggregate(agg_fn, None, filter, None, None) } +#[pyfunction] +#[pyo3(signature = (sort_expression, percentile, filter=None))] +pub fn percentile_cont( + sort_expression: PySortExpr, + percentile: f64, + filter: Option, +) -> PyDataFusionResult { + let agg_fn = + functions_aggregate::expr_fn::percentile_cont(sort_expression.sort, lit(percentile)); + + add_builder_fns_to_aggregate(agg_fn, None, filter, None, None) +} + // We handle last_value explicitly because the signature expects an order_by // https://github.com/apache/datafusion/issues/12376 #[pyfunction] @@ -1031,6 +1045,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(approx_median))?; m.add_wrapped(wrap_pyfunction!(approx_percentile_cont))?; m.add_wrapped(wrap_pyfunction!(approx_percentile_cont_with_weight))?; + m.add_wrapped(wrap_pyfunction!(percentile_cont))?; m.add_wrapped(wrap_pyfunction!(range))?; m.add_wrapped(wrap_pyfunction!(array_agg))?; m.add_wrapped(wrap_pyfunction!(arrow_typeof))?; @@ -1080,7 +1095,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(from_unixtime))?; m.add_wrapped(wrap_pyfunction!(gcd))?; m.add_wrapped(wrap_pyfunction!(greatest))?; - // m.add_wrapped(wrap_pyfunction!(grouping))?; + m.add_wrapped(wrap_pyfunction!(grouping))?; m.add_wrapped(wrap_pyfunction!(in_list))?; m.add_wrapped(wrap_pyfunction!(initcap))?; m.add_wrapped(wrap_pyfunction!(isnan))?; diff --git a/docs/source/user-guide/common-operations/aggregations.rst b/docs/source/user-guide/common-operations/aggregations.rst index e458e5fcb..de24a2ba5 100644 --- a/docs/source/user-guide/common-operations/aggregations.rst +++ b/docs/source/user-guide/common-operations/aggregations.rst @@ -163,6 +163,168 @@ 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")]) +Grouping Sets +------------- + +The default style of aggregation produces one row per group. Sometimes you want a single query to +produce rows at multiple levels of detail — for example, totals per type *and* an overall grand +total, or subtotals for every combination of two columns plus the individual column totals. Writing +separate queries and concatenating them is tedious and runs the data multiple times. Grouping sets +solve this by letting you specify several grouping levels in one pass. + +DataFusion supports three grouping set styles through the +:py:class:`~datafusion.expr.GroupingSet` class: + +- :py:meth:`~datafusion.expr.GroupingSet.rollup` — hierarchical subtotals, like a drill-down report +- :py:meth:`~datafusion.expr.GroupingSet.cube` — every possible subtotal combination, like a pivot table +- :py:meth:`~datafusion.expr.GroupingSet.grouping_sets` — explicitly list exactly which grouping levels you want + +Because result rows come from different grouping levels, a column that is *not* part of a +particular level will be ``null`` in that row. Use :py:func:`~datafusion.functions.grouping` to +distinguish a real ``null`` in the data from one that means "this column was aggregated across." +It returns ``0`` when the column is a grouping key for that row, and ``1`` when it is not. + +Rollup +^^^^^^ + +:py:meth:`~datafusion.expr.GroupingSet.rollup` creates a hierarchy. ``rollup(a, b)`` produces +grouping sets ``(a, b)``, ``(a)``, and ``()`` — like nested subtotals in a report. This is useful +when your columns have a natural hierarchy, such as region → city or type → subtype. + +Suppose we want to summarize Pokemon stats by ``Type 1`` with subtotals and a grand total. With +the default aggregation style we would need two separate queries. With ``rollup`` we get it all at +once: + +.. ipython:: python + + from datafusion.expr import GroupingSet + + df.aggregate( + [GroupingSet.rollup(col_type_1)], + [f.count(col_speed).alias("Count"), + f.avg(col_speed).alias("Avg Speed"), + f.max(col_speed).alias("Max Speed")] + ).sort(col_type_1.sort(ascending=True, nulls_first=True)) + +The first row — where ``Type 1`` is ``null`` — is the grand total across all types. But how do you +tell a grand-total ``null`` apart from a Pokemon that genuinely has no type? The +:py:func:`~datafusion.functions.grouping` function returns ``0`` when the column is a grouping key +for that row and ``1`` when it is aggregated across. + +.. note:: + + Due to an upstream DataFusion limitation + (`apache/datafusion#21411 `_), + ``.alias()`` cannot be applied directly to a ``grouping()`` expression — it will raise an + error at execution time. Instead, use + :py:meth:`~datafusion.dataframe.DataFrame.with_column_renamed` on the result DataFrame to + give the column a readable name. Once the upstream issue is resolved, you will be able to + use ``.alias()`` directly and the workaround below will no longer be necessary. + +The raw column name generated by ``grouping()`` contains internal identifiers, so we use +:py:meth:`~datafusion.dataframe.DataFrame.with_column_renamed` to clean it up: + +.. ipython:: python + + result = df.aggregate( + [GroupingSet.rollup(col_type_1)], + [f.count(col_speed).alias("Count"), + f.avg(col_speed).alias("Avg Speed"), + f.grouping(col_type_1)] + ) + for field in result.schema(): + if field.name.startswith("grouping("): + result = result.with_column_renamed(field.name, "Is Total") + result.sort(col_type_1.sort(ascending=True, nulls_first=True)) + +With two columns the hierarchy becomes more apparent. ``rollup(Type 1, Type 2)`` produces: + +- one row per ``(Type 1, Type 2)`` pair — the most detailed level +- one row per ``Type 1`` — subtotals +- one grand total row + +.. ipython:: python + + df.aggregate( + [GroupingSet.rollup(col_type_1, col_type_2)], + [f.count(col_speed).alias("Count"), + f.avg(col_speed).alias("Avg Speed")] + ).sort( + col_type_1.sort(ascending=True, nulls_first=True), + col_type_2.sort(ascending=True, nulls_first=True) + ) + +Cube +^^^^ + +:py:meth:`~datafusion.expr.GroupingSet.cube` produces every possible subset. ``cube(a, b)`` +produces grouping sets ``(a, b)``, ``(a)``, ``(b)``, and ``()`` — one more than ``rollup`` because +it also includes ``(b)`` alone. This is useful when neither column is "above" the other in a +hierarchy and you want all cross-tabulations. + +For our Pokemon data, ``cube(Type 1, Type 2)`` gives us stats broken down by the type pair, +by ``Type 1`` alone, by ``Type 2`` alone, and a grand total — all in one query: + +.. ipython:: python + + df.aggregate( + [GroupingSet.cube(col_type_1, col_type_2)], + [f.count(col_speed).alias("Count"), + f.avg(col_speed).alias("Avg Speed")] + ).sort( + col_type_1.sort(ascending=True, nulls_first=True), + col_type_2.sort(ascending=True, nulls_first=True) + ) + +Compared to the ``rollup`` example above, notice the extra rows where ``Type 1`` is ``null`` but +``Type 2`` has a value — those are the per-``Type 2`` subtotals that ``rollup`` does not include. + +Explicit Grouping Sets +^^^^^^^^^^^^^^^^^^^^^^ + +:py:meth:`~datafusion.expr.GroupingSet.grouping_sets` lets you list exactly which grouping levels +you need when ``rollup`` or ``cube`` would produce too many or too few. Each argument is a list of +columns forming one grouping set. + +For example, if we want only the per-``Type 1`` totals and per-``Type 2`` totals — but *not* the +full ``(Type 1, Type 2)`` detail rows or the grand total — we can ask for exactly that: + +.. ipython:: python + + df.aggregate( + [GroupingSet.grouping_sets([col_type_1], [col_type_2])], + [f.count(col_speed).alias("Count"), + f.avg(col_speed).alias("Avg Speed")] + ).sort( + col_type_1.sort(ascending=True, nulls_first=True), + col_type_2.sort(ascending=True, nulls_first=True) + ) + +Each row belongs to exactly one grouping level. The :py:func:`~datafusion.functions.grouping` +function tells you which level each row comes from: + +.. ipython:: python + + result = df.aggregate( + [GroupingSet.grouping_sets([col_type_1], [col_type_2])], + [f.count(col_speed).alias("Count"), + f.avg(col_speed).alias("Avg Speed"), + f.grouping(col_type_1), + f.grouping(col_type_2)] + ) + for field in result.schema(): + if field.name.startswith("grouping("): + clean = field.name.split(".")[-1].rstrip(")") + result = result.with_column_renamed(field.name, f"grouping({clean})") + result.sort( + col_type_1.sort(ascending=True, nulls_first=True), + col_type_2.sort(ascending=True, nulls_first=True) + ) + +Where ``grouping(Type 1)`` is ``0`` the row is a per-``Type 1`` total (and ``Type 2`` is ``null``). +Where ``grouping(Type 2)`` is ``0`` the row is a per-``Type 2`` total (and ``Type 1`` is ``null``). + + Aggregate Functions ------------------- @@ -192,6 +354,7 @@ The available aggregate functions are: - :py:func:`datafusion.functions.stddev_pop` - :py:func:`datafusion.functions.var_samp` - :py:func:`datafusion.functions.var_pop` + - :py:func:`datafusion.functions.var_population` 6. Linear Regression Functions - :py:func:`datafusion.functions.regr_count` - :py:func:`datafusion.functions.regr_slope` @@ -208,9 +371,16 @@ The available aggregate functions are: - :py:func:`datafusion.functions.nth_value` 8. String Functions - :py:func:`datafusion.functions.string_agg` -9. Approximation Functions +9. Percentile Functions + - :py:func:`datafusion.functions.percentile_cont` + - :py:func:`datafusion.functions.quantile_cont` - :py:func:`datafusion.functions.approx_distinct` - :py:func:`datafusion.functions.approx_median` - :py:func:`datafusion.functions.approx_percentile_cont` - :py:func:`datafusion.functions.approx_percentile_cont_with_weight` +10. Grouping Set Functions + - :py:func:`datafusion.functions.grouping` + - :py:meth:`datafusion.expr.GroupingSet.rollup` + - :py:meth:`datafusion.expr.GroupingSet.cube` + - :py:meth:`datafusion.expr.GroupingSet.grouping_sets` diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index 10e2a913f..9907eae8b 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -633,8 +633,22 @@ def aggregate( ) -> DataFrame: """Aggregates the rows of the current DataFrame. + By default each unique combination of the ``group_by`` columns + produces one row. To get multiple levels of subtotals in a + single pass, pass a + :py:class:`~datafusion.expr.GroupingSet` expression + (created via + :py:meth:`~datafusion.expr.GroupingSet.rollup`, + :py:meth:`~datafusion.expr.GroupingSet.cube`, or + :py:meth:`~datafusion.expr.GroupingSet.grouping_sets`) + as the ``group_by`` argument. See the + :ref:`aggregation` user guide for detailed examples. + Args: - group_by: Sequence of expressions or column names to group by. + group_by: Sequence of expressions or column names to group + by. A :py:class:`~datafusion.expr.GroupingSet` + expression may be included to produce multiple grouping + levels (rollup, cube, or explicit grouping sets). aggs: Sequence of expressions to aggregate. Returns: diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 14753a4f5..35388468c 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -91,7 +91,6 @@ Extension = expr_internal.Extension FileType = expr_internal.FileType Filter = expr_internal.Filter -GroupingSet = expr_internal.GroupingSet Join = expr_internal.Join ILike = expr_internal.ILike InList = expr_internal.InList @@ -1430,3 +1429,129 @@ def __repr__(self) -> str: SortKey = Expr | SortExpr | str + + +class GroupingSet: + """Factory for creating grouping set expressions. + + Grouping sets control how + :py:meth:`~datafusion.dataframe.DataFrame.aggregate` groups rows. + Instead of a single ``GROUP BY``, they produce multiple grouping + levels in one pass — subtotals, cross-tabulations, or arbitrary + column subsets. + + Use :py:func:`~datafusion.functions.grouping` in the aggregate list + to tell which columns are aggregated across in each result row. + """ + + @staticmethod + def rollup(*exprs: Expr | str) -> Expr: + """Create a ``ROLLUP`` grouping set for use with ``aggregate()``. + + ``ROLLUP`` generates all prefixes of the given column list as + grouping sets. For example, ``rollup(a, b)`` produces grouping + sets ``(a, b)``, ``(a)``, and ``()`` (grand total). + + This is equivalent to ``GROUP BY ROLLUP(a, b)`` in SQL. + + Args: + *exprs: Column expressions or column name strings to + include in the rollup. + + Examples: + >>> from datafusion.expr import GroupingSet + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 1, 2], "b": [10, 20, 30]}) + >>> result = df.aggregate( + ... [GroupingSet.rollup(dfn.col("a"))], + ... [dfn.functions.sum(dfn.col("b")).alias("s"), + ... dfn.functions.grouping(dfn.col("a"))], + ... ).sort(dfn.col("a").sort(nulls_first=False)) + >>> result.collect_column("s").to_pylist() + [30, 30, 60] + + See Also: + :py:meth:`cube`, :py:meth:`grouping_sets`, + :py:func:`~datafusion.functions.grouping` + """ + args = [_to_raw_expr(e) for e in exprs] + return Expr(expr_internal.GroupingSet.rollup(*args)) + + @staticmethod + def cube(*exprs: Expr | str) -> Expr: + """Create a ``CUBE`` grouping set for use with ``aggregate()``. + + ``CUBE`` generates all possible subsets of the given column list + as grouping sets. For example, ``cube(a, b)`` produces grouping + sets ``(a, b)``, ``(a)``, ``(b)``, and ``()`` (grand total). + + This is equivalent to ``GROUP BY CUBE(a, b)`` in SQL. + + Args: + *exprs: Column expressions or column name strings to + include in the cube. + + Examples: + With a single column, ``cube`` behaves identically to + :py:meth:`rollup`: + + >>> from datafusion.expr import GroupingSet + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 1, 2], "b": [10, 20, 30]}) + >>> result = df.aggregate( + ... [GroupingSet.cube(dfn.col("a"))], + ... [dfn.functions.sum(dfn.col("b")).alias("s"), + ... dfn.functions.grouping(dfn.col("a"))], + ... ).sort(dfn.col("a").sort(nulls_first=False)) + >>> result.collect_column("s").to_pylist() + [30, 30, 60] + + See Also: + :py:meth:`rollup`, :py:meth:`grouping_sets`, + :py:func:`~datafusion.functions.grouping` + """ + args = [_to_raw_expr(e) for e in exprs] + return Expr(expr_internal.GroupingSet.cube(*args)) + + @staticmethod + def grouping_sets(*expr_lists: list[Expr | str]) -> Expr: + """Create explicit grouping sets for use with ``aggregate()``. + + Each argument is a list of column expressions or column name + strings representing one grouping set. For example, + ``grouping_sets([a], [b])`` groups by ``a`` alone and by ``b`` + alone in a single query. + + This is equivalent to ``GROUP BY GROUPING SETS ((a), (b))`` in + SQL. + + Args: + *expr_lists: Each positional argument is a list of + expressions or column name strings forming one + grouping set. + + Examples: + >>> from datafusion.expr import GroupingSet + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict( + ... {"a": ["x", "x", "y"], "b": ["m", "n", "m"], + ... "c": [1, 2, 3]}) + >>> result = df.aggregate( + ... [GroupingSet.grouping_sets( + ... [dfn.col("a")], [dfn.col("b")])], + ... [dfn.functions.sum(dfn.col("c")).alias("s"), + ... dfn.functions.grouping(dfn.col("a")), + ... dfn.functions.grouping(dfn.col("b"))], + ... ).sort( + ... dfn.col("a").sort(nulls_first=False), + ... dfn.col("b").sort(nulls_first=False), + ... ) + >>> result.collect_column("s").to_pylist() + [3, 3, 4, 2] + + See Also: + :py:meth:`rollup`, :py:meth:`cube`, + :py:func:`~datafusion.functions.grouping` + """ + raw_lists = [[_to_raw_expr(e) for e in lst] for lst in expr_lists] + return Expr(expr_internal.GroupingSet.grouping_sets(*raw_lists)) diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index aa7f28746..9dfabb62d 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -166,6 +166,7 @@ "generate_series", "get_field", "greatest", + "grouping", "ifnull", "in_list", "initcap", @@ -256,9 +257,11 @@ "order_by", "overlay", "percent_rank", + "percentile_cont", "pi", "pow", "power", + "quantile_cont", "radians", "random", "range", @@ -331,6 +334,7 @@ "uuid", "var", "var_pop", + "var_population", "var_samp", "var_sample", "version", @@ -2654,7 +2658,6 @@ def arrow_cast(expr: Expr, data_type: Expr | str | pa.DataType) -> Expr: >>> result.collect_column("c")[0].as_py() 1.0 - >>> import pyarrow as pa >>> result = df.select( ... dfn.functions.arrow_cast( ... dfn.col("a"), data_type=pa.float64() @@ -2677,7 +2680,6 @@ def arrow_metadata(expr: Expr, key: Expr | str | None = None) -> Expr: If called with two arguments, returns the value for the specified metadata key. Examples: - >>> import pyarrow as pa >>> field = pa.field("val", pa.int64(), metadata={"k": "v"}) >>> schema = pa.schema([field]) >>> batch = pa.RecordBatch.from_arrays([pa.array([1])], schema=schema) @@ -2746,7 +2748,6 @@ def union_extract(union_expr: Expr, field_name: Expr | str) -> Expr: variant, otherwise returns NULL. Examples: - >>> import pyarrow as pa >>> ctx = dfn.SessionContext() >>> types = pa.array([0, 1, 0], type=pa.int8()) >>> offsets = pa.array([0, 0, 1], type=pa.int32()) @@ -2771,7 +2772,6 @@ def union_tag(union_expr: Expr) -> Expr: """Returns the tag (active field name) of a union type. Examples: - >>> import pyarrow as pa >>> ctx = dfn.SessionContext() >>> types = pa.array([0, 1, 0], type=pa.int8()) >>> offsets = pa.array([0, 0, 1], type=pa.int32()) @@ -4306,6 +4306,60 @@ def approx_percentile_cont_with_weight( ) +def percentile_cont( + sort_expression: Expr | SortExpr, + percentile: float, + filter: Expr | None = None, +) -> Expr: + """Computes the exact percentile of input values using continuous interpolation. + + Unlike :py:func:`approx_percentile_cont`, this function computes the exact + percentile value rather than an approximation. + + If using the builder functions described in ref:`_aggregation` this function ignores + the options ``order_by``, ``null_treatment``, and ``distinct``. + + Args: + sort_expression: Values for which to find the percentile + percentile: This must be between 0.0 and 1.0, inclusive + filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.percentile_cont( + ... dfn.col("a"), 0.5 + ... ).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3.0 + + >>> result = df.aggregate( + ... [], [dfn.functions.percentile_cont( + ... dfn.col("a"), 0.5, + ... filter=dfn.col("a") > dfn.lit(1.0), + ... ).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3.5 + """ + sort_expr_raw = sort_or_default(sort_expression) + filter_raw = filter.expr if filter is not None else None + return Expr(f.percentile_cont(sort_expr_raw, percentile, filter=filter_raw)) + + +def quantile_cont( + sort_expression: Expr | SortExpr, + percentile: float, + filter: Expr | None = None, +) -> Expr: + """Computes the exact percentile of input values using continuous interpolation. + + See Also: + This is an alias for :py:func:`percentile_cont`. + """ + return percentile_cont(sort_expression, percentile, filter) + + def array_agg( expression: Expr, distinct: bool = False, @@ -4364,6 +4418,65 @@ def array_agg( ) +def grouping( + expression: Expr, + distinct: bool = False, + filter: Expr | None = None, +) -> Expr: + """Indicates whether a column is aggregated across in the current row. + + Returns 0 when the column is part of the grouping key for that row + (i.e., the row contains per-group results for that column). Returns 1 + when the column is *not* part of the grouping key (i.e., the row's + aggregate spans all values of that column). + + This function is meaningful with + :py:meth:`GroupingSet.rollup `, + :py:meth:`GroupingSet.cube `, or + :py:meth:`GroupingSet.grouping_sets `, + where different rows are grouped by different subsets of columns. In a + default aggregation without grouping sets every column is always part + of the key, so ``grouping()`` always returns 0. + + .. warning:: + + Due to an upstream DataFusion limitation + (`#21411 `_), + ``.alias()`` cannot be applied directly to a ``grouping()`` + expression. Doing so will raise an error at execution time. To + rename the column, use + :py:meth:`~datafusion.dataframe.DataFrame.with_column_renamed` + on the result DataFrame instead. + + Args: + expression: The column to check grouping status for + distinct: If True, compute on distinct values only + filter: If provided, only compute against rows for which the filter is True + + Examples: + With :py:meth:`~datafusion.expr.GroupingSet.rollup`, the result + includes both per-group rows (``grouping(a) = 0``) and a + grand-total row where ``a`` is aggregated across + (``grouping(a) = 1``): + + >>> from datafusion.expr import GroupingSet + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 1, 2], "b": [10, 20, 30]}) + >>> result = df.aggregate( + ... [GroupingSet.rollup(dfn.col("a"))], + ... [dfn.functions.sum(dfn.col("b")).alias("s"), + ... dfn.functions.grouping(dfn.col("a"))], + ... ).sort(dfn.col("a").sort(nulls_first=False)) + >>> result.collect_column("s").to_pylist() + [30, 30, 60] + + See Also: + :py:class:`~datafusion.expr.GroupingSet` + """ + filter_raw = filter.expr if filter is not None else None + return Expr(f.grouping(expression.expr, distinct=distinct, filter=filter_raw)) + + def avg( expression: Expr, filter: Expr | None = None, @@ -4835,6 +4948,15 @@ def var_pop(expression: Expr, filter: Expr | None = None) -> Expr: return Expr(f.var_pop(expression.expr, filter=filter_raw)) +def var_population(expression: Expr, filter: Expr | None = None) -> Expr: + """Computes the population variance of the argument. + + See Also: + This is an alias for :py:func:`var_pop`. + """ + return var_pop(expression, filter) + + def var_samp(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the sample variance of the argument. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 4e99fa9e3..11e94af1c 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -22,6 +22,7 @@ import pytest from datafusion import SessionContext, column, literal from datafusion import functions as f +from datafusion.expr import GroupingSet np.seterr(invalid="ignore") @@ -1820,6 +1821,114 @@ def test_conditional_functions(df_with_nulls, expr, expected): assert result.column(0) == expected +@pytest.mark.parametrize( + ("func", "filter_expr", "expected"), + [ + (f.percentile_cont, None, 3.0), + (f.percentile_cont, column("a") > literal(1.0), 3.5), + (f.quantile_cont, None, 3.0), + ], + ids=["no_filter", "with_filter", "quantile_cont_alias"], +) +def test_percentile_cont(func, filter_expr, expected): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}) + result = df.aggregate( + [], [func(column("a"), 0.5, filter=filter_expr).alias("v")] + ).collect()[0] + assert result.column(0)[0].as_py() == expected + + +@pytest.mark.parametrize( + ("grouping_set_expr", "expected_grouping", "expected_sums"), + [ + (GroupingSet.rollup(column("a")), [0, 0, 1], [30, 30, 60]), + (GroupingSet.cube(column("a")), [0, 0, 1], [30, 30, 60]), + (GroupingSet.rollup("a"), [0, 0, 1], [30, 30, 60]), + (GroupingSet.cube("a"), [0, 0, 1], [30, 30, 60]), + ], + ids=["rollup", "cube", "rollup_str", "cube_str"], +) +def test_grouping_set_single_column( + grouping_set_expr, expected_grouping, expected_sums +): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1, 1, 2], "b": [10, 20, 30]}) + result = df.aggregate( + [grouping_set_expr], + [f.sum(column("b")).alias("s"), f.grouping(column("a"))], + ).sort(column("a").sort(ascending=True, nulls_first=False)) + batches = result.collect() + g = pa.concat_arrays([b.column(2) for b in batches]).to_pylist() + s = pa.concat_arrays([b.column("s") for b in batches]).to_pylist() + assert g == expected_grouping + assert s == expected_sums + + +@pytest.mark.parametrize( + ("grouping_set_expr", "expected_rows"), + [ + # rollup(a, b) => (a,b), (a), () => 3 + 2 + 1 = 6 + (GroupingSet.rollup(column("a"), column("b")), 6), + # cube(a, b) => (a,b), (a), (b), () => 3 + 2 + 2 + 1 = 8 + (GroupingSet.cube(column("a"), column("b")), 8), + (GroupingSet.rollup("a", "b"), 6), + (GroupingSet.cube("a", "b"), 8), + ], + ids=["rollup", "cube", "rollup_str", "cube_str"], +) +def test_grouping_set_multi_column(grouping_set_expr, expected_rows): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1, 1, 2], "b": ["x", "y", "x"], "c": [10, 20, 30]}) + result = df.aggregate( + [grouping_set_expr], + [f.sum(column("c")).alias("s")], + ) + total_rows = sum(b.num_rows for b in result.collect()) + assert total_rows == expected_rows + + +@pytest.mark.parametrize( + "grouping_set_expr", + [ + GroupingSet.grouping_sets([column("a")], [column("b")]), + GroupingSet.grouping_sets(["a"], ["b"]), + ], + ids=["expr", "str"], +) +def test_grouping_sets_explicit(grouping_set_expr): + # Each row's grouping() value tells you which columns are aggregated across. + ctx = SessionContext() + df = ctx.from_pydict({"a": ["x", "x", "y"], "b": ["m", "n", "m"], "c": [1, 2, 3]}) + result = df.aggregate( + [grouping_set_expr], + [ + f.sum(column("c")).alias("s"), + f.grouping(column("a")), + f.grouping(column("b")), + ], + ).sort( + column("a").sort(ascending=True, nulls_first=False), + column("b").sort(ascending=True, nulls_first=False), + ) + batches = result.collect() + ga = pa.concat_arrays([b.column(3) for b in batches]).to_pylist() + gb = pa.concat_arrays([b.column(4) for b in batches]).to_pylist() + # Rows grouped by (a): ga=0 (a is a key), gb=1 (b is aggregated across) + # Rows grouped by (b): ga=1 (a is aggregated across), gb=0 (b is a key) + assert ga == [0, 0, 1, 1] + assert gb == [1, 1, 0, 0] + + +def test_var_population(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [-1.0, 0.0, 2.0]}) + result = df.aggregate([], [f.var_population(column("a")).alias("v")]).collect()[0] + # var_population is an alias for var_pop + expected = df.aggregate([], [f.var_pop(column("a")).alias("v")]).collect()[0] + assert abs(result.column(0)[0].as_py() - expected.column(0)[0].as_py()) < 1e-10 + + def test_get_field(df): df = df.with_column( "s", From 52932128d353e417ddae2c5ff3f14135cb806f7e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 7 Apr 2026 14:58:09 -0400 Subject: [PATCH 13/83] Add missing Dataframe functions (#1472) * Add missing DataFrame methods for set operations and query Expose upstream DataFusion DataFrame methods that were not yet available in the Python API. Closes #1455. Set operations: - except_distinct: set difference with deduplication - intersect_distinct: set intersection with deduplication - union_by_name: union matching columns by name instead of position - union_by_name_distinct: union by name with deduplication Query: - distinct_on: deduplicate rows based on specific columns - sort_by: sort by expressions with ascending order and nulls last Note: show_limit is already covered by the existing show(num) method. explain_with_options and with_param_values are deferred as they require exposing additional types (ExplainOption, ParamValues). Co-Authored-By: Claude Opus 4.6 (1M context) * Add ExplainFormat enum and format option to DataFrame.explain() Extend the existing explain() method with an optional format parameter instead of adding a separate explain_with_options() method. This keeps the API simple while exposing all upstream ExplainOption functionality. Available formats: indent (default), tree, pgjson, graphviz. The ExplainFormat enum is exported from the top-level datafusion module. Co-Authored-By: Claude Opus 4.6 (1M context) * Add DataFrame.window() and unnest recursion options Expose remaining DataFrame methods from upstream DataFusion. Closes #1456. - window(*exprs): apply window function expressions and append results as new columns - unnest_column/unnest_columns: add optional recursions parameter for controlling unnest depth via (input_column, output_column, depth) tuples Note: drop_columns is already exposed as the existing drop() method. Co-Authored-By: Claude Opus 4.6 (1M context) * Update docstring Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Improve docstrings and test robustness for new DataFrame methods Clarify except_distinct/intersect_distinct docstrings, add deterministic sort to test_window, add sort_by ascending verification test, and add smoke tests for PGJSON and GRAPHVIZ explain formats. Co-Authored-By: Claude Opus 4.6 (1M context) * Consolidate new DataFrame tests into parametrized tests Combine set operation tests (except_distinct, intersect_distinct, union_by_name, union_by_name_distinct) into a single parametrized test_set_operations_distinct. Merge sort_by tests and convert explain format tests to parametrized form. Co-Authored-By: Claude Opus 4.6 (1M context) * Add doctest examples to new DataFrame method docstrings Add >>> style usage examples for window, explain, except_distinct, intersect_distinct, union_by_name, union_by_name_distinct, distinct_on, sort_by, and unnest_columns to match existing docstring conventions. Co-Authored-By: Claude Opus 4.6 (1M context) * Improve error messages, tests, and API hygiene from PR review - Provide actionable error message for invalid explain format strings - Remove recursions param from deprecated unnest_column (use unnest_columns) - Add null-handling test case for sort_by to verify nulls-last behavior - Add format-specific assertions to explain tests (TREE, PGJSON, GRAPHVIZ) - Add deep recursion test for unnest_columns with depth > 1 - Add multi-expression window test to verify variadic *exprs Co-Authored-By: Claude Opus 4.6 (1M context) * Consolidate window and unnest tests into parametrized tests Combine test_window and test_window_multiple_expressions into a single parametrized test. Merge unnest recursion tests into one parametrized test covering basic, explicit depth 1, and deep recursion cases. Co-Authored-By: Claude Opus 4.6 (1M context) * Address PR review feedback for DataFrame operations - Use upstream parse error for explain format instead of hardcoded options - Fix sort_by to use column name resolution consistent with sort() - Use ExplainFormat enum members directly in tests instead of string lookup - Merge union_by_name_distinct into union_by_name(distinct=False) for a more Pythonic API - Update check-upstream skill to note union_by_name_distinct coverage Co-Authored-By: Claude Opus 4.6 (1M context) * Add DataFrame.column(), col(), and find_qualified_columns() methods Expose upstream find_qualified_columns to resolve unqualified column names into fully qualified column expressions. This is especially useful for disambiguating columns after joins. - find_qualified_columns(*names) on Rust side calls upstream directly - DataFrame.column(name) and col(name) alias on Python side - Update join and join_on docstrings to reference DataFrame.col() - Add "Disambiguating Columns with DataFrame.col()" section to joins docs - Add tests for qualified column resolution, ambiguity, and join usage Co-Authored-By: Claude Opus 4.6 (1M context) * Merge union_by_name and union_by_name_distinct into a single method with distinct flag Co-Authored-By: Claude Opus 4.6 (1M context) * converting into a python dict loses a column when the names are identical * Consolidate except_all/except_distinct and intersect/intersect_distinct into single methods with distinct flag Follows the same pattern as union(distinct=) and union_by_name(distinct=). Also deprecates union_distinct() in favor of union(distinct=True). Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .ai/skills/check-upstream/SKILL.md | 1 + crates/core/src/dataframe.rs | 157 ++++++-- .../user-guide/common-operations/joins.rst | 33 ++ python/datafusion/__init__.py | 2 + python/datafusion/dataframe.py | 365 +++++++++++++++++- python/tests/test_dataframe.py | 261 +++++++++++++ 6 files changed, 767 insertions(+), 52 deletions(-) diff --git a/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md index f77210371..ac4835a4e 100644 --- a/.ai/skills/check-upstream/SKILL.md +++ b/.ai/skills/check-upstream/SKILL.md @@ -109,6 +109,7 @@ The user may specify an area via `$ARGUMENTS`. If no area is specified or "all" **Evaluated and not requiring separate Python exposure:** - `show_limit` — already covered by `DataFrame.show()`, which provides the same functionality with a simpler API - `with_param_values` — already covered by the `param_values` argument on `SessionContext.sql()`, which accomplishes the same thing more robustly +- `union_by_name_distinct` — already covered by `DataFrame.union_by_name(distinct=True)`, which provides a more Pythonic API **How to check:** 1. Fetch the upstream DataFrame documentation page listing all methods diff --git a/crates/core/src/dataframe.rs b/crates/core/src/dataframe.rs index 72595ba81..fff5118d5 100644 --- a/crates/core/src/dataframe.rs +++ b/crates/core/src/dataframe.rs @@ -582,6 +582,14 @@ impl PyDataFrame { Ok(Self::new(df)) } + /// Apply window function expressions to the DataFrame + #[pyo3(signature = (*exprs))] + fn window(&self, exprs: Vec) -> PyDataFusionResult { + let window_exprs = exprs.into_iter().map(|e| e.into()).collect(); + let df = self.df.as_ref().clone().window(window_exprs)?; + Ok(Self::new(df)) + } + fn filter(&self, predicate: PyExpr) -> PyDataFusionResult { let df = self.df.as_ref().clone().filter(predicate.into())?; Ok(Self::new(df)) @@ -804,9 +812,27 @@ impl PyDataFrame { } /// Print the query plan - #[pyo3(signature = (verbose=false, analyze=false))] - fn explain(&self, py: Python, verbose: bool, analyze: bool) -> PyDataFusionResult<()> { - let df = self.df.as_ref().clone().explain(verbose, analyze)?; + #[pyo3(signature = (verbose=false, analyze=false, format=None))] + fn explain( + &self, + py: Python, + verbose: bool, + analyze: bool, + format: Option<&str>, + ) -> PyDataFusionResult<()> { + let explain_format = match format { + Some(f) => f + .parse::() + .map_err(|e| { + PyDataFusionError::Common(format!("Invalid explain format '{}': {}", f, e)) + })?, + None => datafusion::common::format::ExplainFormat::Indent, + }; + let opts = datafusion::logical_expr::ExplainOption::default() + .with_verbose(verbose) + .with_analyze(analyze) + .with_format(explain_format); + let df = self.df.as_ref().clone().explain_with_options(opts)?; print_dataframe(py, df) } @@ -864,22 +890,14 @@ impl PyDataFrame { Ok(Self::new(new_df)) } - /// Calculate the distinct union of two `DataFrame`s. The - /// two `DataFrame`s must have exactly the same schema - fn union_distinct(&self, py_df: PyDataFrame) -> PyDataFusionResult { - let new_df = self - .df - .as_ref() - .clone() - .union_distinct(py_df.df.as_ref().clone())?; - Ok(Self::new(new_df)) - } - - #[pyo3(signature = (column, preserve_nulls=true))] - fn unnest_column(&self, column: &str, preserve_nulls: bool) -> PyDataFusionResult { - // TODO: expose RecursionUnnestOptions - // REF: https://github.com/apache/datafusion/pull/11577 - let unnest_options = UnnestOptions::default().with_preserve_nulls(preserve_nulls); + #[pyo3(signature = (column, preserve_nulls=true, recursions=None))] + fn unnest_column( + &self, + column: &str, + preserve_nulls: bool, + recursions: Option>, + ) -> PyDataFusionResult { + let unnest_options = build_unnest_options(preserve_nulls, recursions); let df = self .df .as_ref() @@ -888,15 +906,14 @@ impl PyDataFrame { Ok(Self::new(df)) } - #[pyo3(signature = (columns, preserve_nulls=true))] + #[pyo3(signature = (columns, preserve_nulls=true, recursions=None))] fn unnest_columns( &self, columns: Vec, preserve_nulls: bool, + recursions: Option>, ) -> PyDataFusionResult { - // TODO: expose RecursionUnnestOptions - // REF: https://github.com/apache/datafusion/pull/11577 - let unnest_options = UnnestOptions::default().with_preserve_nulls(preserve_nulls); + let unnest_options = build_unnest_options(preserve_nulls, recursions); let cols = columns.iter().map(|s| s.as_ref()).collect::>(); let df = self .df @@ -907,21 +924,79 @@ impl PyDataFrame { } /// Calculate the intersection of two `DataFrame`s. The two `DataFrame`s must have exactly the same schema - fn intersect(&self, py_df: PyDataFrame) -> PyDataFusionResult { - let new_df = self - .df - .as_ref() - .clone() - .intersect(py_df.df.as_ref().clone())?; + #[pyo3(signature = (py_df, distinct=false))] + fn intersect(&self, py_df: PyDataFrame, distinct: bool) -> PyDataFusionResult { + let base = self.df.as_ref().clone(); + let other = py_df.df.as_ref().clone(); + let new_df = if distinct { + base.intersect_distinct(other)? + } else { + base.intersect(other)? + }; Ok(Self::new(new_df)) } /// Calculate the exception of two `DataFrame`s. The two `DataFrame`s must have exactly the same schema - fn except_all(&self, py_df: PyDataFrame) -> PyDataFusionResult { - let new_df = self.df.as_ref().clone().except(py_df.df.as_ref().clone())?; + #[pyo3(signature = (py_df, distinct=false))] + fn except_all(&self, py_df: PyDataFrame, distinct: bool) -> PyDataFusionResult { + let base = self.df.as_ref().clone(); + let other = py_df.df.as_ref().clone(); + let new_df = if distinct { + base.except_distinct(other)? + } else { + base.except(other)? + }; Ok(Self::new(new_df)) } + /// Union two DataFrames matching columns by name + #[pyo3(signature = (py_df, distinct=false))] + fn union_by_name(&self, py_df: PyDataFrame, distinct: bool) -> PyDataFusionResult { + let base = self.df.as_ref().clone(); + let other = py_df.df.as_ref().clone(); + let new_df = if distinct { + base.union_by_name_distinct(other)? + } else { + base.union_by_name(other)? + }; + Ok(Self::new(new_df)) + } + + /// Deduplicate rows based on specific columns, keeping the first row per group + fn distinct_on( + &self, + on_expr: Vec, + select_expr: Vec, + sort_expr: Option>, + ) -> PyDataFusionResult { + let on_expr = on_expr.into_iter().map(|e| e.into()).collect(); + let select_expr = select_expr.into_iter().map(|e| e.into()).collect(); + let sort_expr = sort_expr.map(to_sort_expressions); + let df = self + .df + .as_ref() + .clone() + .distinct_on(on_expr, select_expr, sort_expr)?; + Ok(Self::new(df)) + } + + /// Sort by column expressions with ascending order and nulls last + fn sort_by(&self, exprs: Vec) -> PyDataFusionResult { + let exprs = exprs.into_iter().map(|e| e.into()).collect(); + let df = self.df.as_ref().clone().sort_by(exprs)?; + Ok(Self::new(df)) + } + + /// Return fully qualified column expressions for the given column names + fn find_qualified_columns(&self, names: Vec) -> PyDataFusionResult> { + let name_refs: Vec<&str> = names.iter().map(|s| s.as_str()).collect(); + let qualified = self.df.find_qualified_columns(&name_refs)?; + Ok(qualified + .into_iter() + .map(|q| Expr::Column(Column::from(q)).into()) + .collect()) + } + /// Write a `DataFrame` to a CSV file. fn write_csv( &self, @@ -1295,6 +1370,26 @@ impl PyDataFrameWriteOptions { } } +fn build_unnest_options( + preserve_nulls: bool, + recursions: Option>, +) -> UnnestOptions { + let mut opts = UnnestOptions::default().with_preserve_nulls(preserve_nulls); + if let Some(recs) = recursions { + opts.recursions = recs + .into_iter() + .map( + |(input, output, depth)| datafusion::common::RecursionUnnestOption { + input_column: datafusion::common::Column::from(input.as_str()), + output_column: datafusion::common::Column::from(output.as_str()), + depth, + }, + ) + .collect(); + } + opts +} + /// Print DataFrame fn print_dataframe(py: Python, df: DataFrame) -> PyDataFusionResult<()> { // Get string representation of record batches diff --git a/docs/source/user-guide/common-operations/joins.rst b/docs/source/user-guide/common-operations/joins.rst index 1d9d70385..a289c9377 100644 --- a/docs/source/user-guide/common-operations/joins.rst +++ b/docs/source/user-guide/common-operations/joins.rst @@ -134,3 +134,36 @@ In contrast to the above example, if we wish to get both columns: .. ipython:: python left.join(right, "id", how="inner", coalesce_duplicate_keys=False) + +Disambiguating Columns with ``DataFrame.col()`` +------------------------------------------------ + +When both DataFrames contain non-key columns with the same name, you can use +:py:meth:`~datafusion.dataframe.DataFrame.col` on each DataFrame **before** the +join to create fully qualified column references. These references can then be +used in the join predicate and when selecting from the result. + +This is especially useful with :py:meth:`~datafusion.dataframe.DataFrame.join_on`, +which accepts expression-based predicates. + +.. ipython:: python + + left = ctx.from_pydict( + { + "id": [1, 2, 3], + "val": [10, 20, 30], + } + ) + + right = ctx.from_pydict( + { + "id": [1, 2, 3], + "val": [40, 50, 60], + } + ) + + joined = left.join_on( + right, left.col("id") == right.col("id"), how="inner" + ) + + joined.select(left.col("id"), left.col("val"), right.col("val")) diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 2e6f81166..a736c3966 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -47,6 +47,7 @@ from .dataframe import ( DataFrame, DataFrameWriteOptions, + ExplainFormat, InsertOp, ParquetColumnOptions, ParquetWriterOptions, @@ -82,6 +83,7 @@ "DataFrameWriteOptions", "Database", "ExecutionPlan", + "ExplainFormat", "Expr", "InsertOp", "LogicalPlan", diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index 9907eae8b..9dc5f0e7d 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -44,6 +44,7 @@ Expr, SortExpr, SortKey, + _to_raw_expr, ensure_expr, ensure_expr_list, expr_list_to_raw_expr_list, @@ -65,6 +66,25 @@ from enum import Enum +class ExplainFormat(Enum): + """Output format for explain plans. + + Controls how the query plan is rendered in :py:meth:`DataFrame.explain`. + """ + + INDENT = "indent" + """Default indented text format.""" + + TREE = "tree" + """Tree-style visual format with box-drawing characters.""" + + PGJSON = "pgjson" + """PostgreSQL-compatible JSON format for use with visualization tools.""" + + GRAPHVIZ = "graphviz" + """Graphviz DOT format for graph rendering.""" + + # excerpt from deltalake # https://github.com/apache/datafusion-python/pull/981#discussion_r1905619163 class Compression(Enum): @@ -395,6 +415,80 @@ def schema(self) -> pa.Schema: """ return self.df.schema() + def column(self, name: str) -> Expr: + """Return a fully qualified column expression for ``name``. + + Resolves an unqualified column name against this DataFrame's schema + and returns an :py:class:`Expr` whose underlying column reference + includes the table qualifier. This is especially useful after joins, + where the same column name may appear in multiple relations. + + Args: + name: Unqualified column name to look up. + + Returns: + A fully qualified column expression. + + Raises: + Exception: If the column is not found or is ambiguous (exists in + multiple relations). + + Examples: + Resolve a column from a simple DataFrame: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2], "b": [3, 4]}) + >>> expr = df.column("a") + >>> df.select(expr).to_pydict() + {'a': [1, 2]} + + Resolve qualified columns after a join: + + >>> left = ctx.from_pydict({"id": [1, 2], "x": [10, 20]}) + >>> right = ctx.from_pydict({"id": [1, 2], "y": [30, 40]}) + >>> joined = left.join(right, on="id", how="inner") + >>> expr = joined.column("y") + >>> joined.select("id", expr).sort("id").to_pydict() + {'id': [1, 2], 'y': [30, 40]} + """ + return self.find_qualified_columns(name)[0] + + def col(self, name: str) -> Expr: + """Alias for :py:meth:`column`. + + See Also: + :py:meth:`column` + """ + return self.column(name) + + def find_qualified_columns(self, *names: str) -> list[Expr]: + """Return fully qualified column expressions for the given names. + + This is a batch version of :py:meth:`column` — it resolves each + unqualified name against the DataFrame's schema and returns a list + of qualified column expressions. + + Args: + names: Unqualified column names to look up. + + Returns: + List of fully qualified column expressions, one per name. + + Raises: + Exception: If any column is not found or is ambiguous. + + Examples: + Resolve multiple columns at once: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2], "b": [3, 4], "c": [5, 6]}) + >>> exprs = df.find_qualified_columns("a", "c") + >>> df.select(*exprs).to_pydict() + {'a': [1, 2], 'c': [5, 6]} + """ + raw_exprs = self.df.find_qualified_columns(list(names)) + return [Expr(e) for e in raw_exprs] + @deprecated( "select_columns() is deprecated. Use :py:meth:`~DataFrame.select` instead" ) @@ -468,6 +562,36 @@ def drop(self, *columns: str) -> DataFrame: """ return DataFrame(self.df.drop(*columns)) + def window(self, *exprs: Expr) -> DataFrame: + """Add window function columns to the DataFrame. + + Applies the given window function expressions and appends the results + as new columns. + + Args: + exprs: Window function expressions to evaluate. + + Returns: + DataFrame with new window function columns appended. + + Examples: + Add a row number within each group: + + >>> import datafusion.functions as f + >>> from datafusion import col + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3], "b": ["x", "x", "y"]}) + >>> df = df.window( + ... f.row_number( + ... partition_by=[col("b")], order_by=[col("a")] + ... ).alias("rn") + ... ) + >>> "rn" in df.schema().names + True + """ + raw = expr_list_to_raw_expr_list(exprs) + return DataFrame(self.df.window(*raw)) + def filter(self, *predicates: Expr | str) -> DataFrame: """Return a DataFrame for which ``predicate`` evaluates to ``True``. @@ -837,7 +961,13 @@ def join( ) -> DataFrame: """Join this :py:class:`DataFrame` with another :py:class:`DataFrame`. - `on` has to be provided or both `left_on` and `right_on` in conjunction. + ``on`` has to be provided or both ``left_on`` and ``right_on`` in + conjunction. + + When non-key columns share the same name in both DataFrames, use + :py:meth:`DataFrame.col` on each DataFrame **before** the join to + obtain fully qualified column references that can disambiguate them. + See :py:meth:`join_on` for an example. Args: right: Other DataFrame to join with. @@ -911,7 +1041,14 @@ def join_on( built with :func:`datafusion.col`. On expressions are used to support in-equality predicates. Equality predicates are correctly optimized. + Use :py:meth:`DataFrame.col` on each DataFrame **before** the join to + obtain fully qualified column references. These qualified references + can then be used in the join predicate and to disambiguate columns + with the same name when selecting from the result. + Examples: + Join with unique column names: + >>> ctx = dfn.SessionContext() >>> left = ctx.from_pydict({"a": [1, 2], "x": ["a", "b"]}) >>> right = ctx.from_pydict({"b": [1, 2], "y": ["c", "d"]}) @@ -920,6 +1057,18 @@ def join_on( ... ).sort(col("x")).to_pydict() {'a': [1, 2], 'x': ['a', 'b'], 'b': [1, 2], 'y': ['c', 'd']} + Use :py:meth:`col` to disambiguate shared column names: + + >>> left = ctx.from_pydict({"id": [1, 2], "val": [10, 20]}) + >>> right = ctx.from_pydict({"id": [1, 2], "val": [30, 40]}) + >>> joined = left.join_on( + ... right, left.col("id") == right.col("id"), how="inner" + ... ) + >>> joined.select( + ... left.col("id"), left.col("val"), right.col("val").alias("rval") + ... ).sort(left.col("id")).to_pydict() + {'id': [1, 2], 'val': [10, 20], 'rval': [30, 40]} + Args: right: Other DataFrame to join with. on_exprs: single or multiple (in)-equality predicates. @@ -932,7 +1081,12 @@ def join_on( exprs = [ensure_expr(expr) for expr in on_exprs] return DataFrame(self.df.join_on(right.df, exprs, how)) - def explain(self, verbose: bool = False, analyze: bool = False) -> None: + def explain( + self, + verbose: bool = False, + analyze: bool = False, + format: ExplainFormat | None = None, + ) -> None: """Print an explanation of the DataFrame's plan so far. If ``analyze`` is specified, runs the plan and reports metrics. @@ -940,8 +1094,23 @@ def explain(self, verbose: bool = False, analyze: bool = False) -> None: Args: verbose: If ``True``, more details will be included. analyze: If ``True``, the plan will run and metrics reported. + format: Output format for the plan. Defaults to + :py:attr:`ExplainFormat.INDENT`. + + Examples: + Show the plan in tree format: + + >>> from datafusion import ExplainFormat + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> df.explain(format=ExplainFormat.TREE) # doctest: +SKIP + + Show plan with runtime metrics: + + >>> df.explain(analyze=True) # doctest: +SKIP """ - self.df.explain(verbose, analyze) + fmt = format.value if format is not None else None + self.df.explain(verbose, analyze, fmt) def logical_plan(self) -> LogicalPlan: """Return the unoptimized ``LogicalPlan``. @@ -1010,45 +1179,170 @@ def union(self, other: DataFrame, distinct: bool = False) -> DataFrame: """ return DataFrame(self.df.union(other.df, distinct)) + @deprecated( + "union_distinct() is deprecated. Use union(other, distinct=True) instead." + ) def union_distinct(self, other: DataFrame) -> DataFrame: """Calculate the distinct union of two :py:class:`DataFrame`. + See Also: + :py:meth:`union` + """ + return self.union(other, distinct=True) + + def intersect(self, other: DataFrame, distinct: bool = False) -> DataFrame: + """Calculate the intersection of two :py:class:`DataFrame`. + The two :py:class:`DataFrame` must have exactly the same schema. - Any duplicate rows are discarded. Args: - other: DataFrame to union with. + other: DataFrame to intersect with. + distinct: If ``True``, duplicate rows are removed from the result. Returns: - DataFrame after union. + DataFrame after intersection. + + Examples: + Find rows common to both DataFrames: + + >>> ctx = dfn.SessionContext() + >>> df1 = ctx.from_pydict({"a": [1, 2, 3], "b": [10, 20, 30]}) + >>> df2 = ctx.from_pydict({"a": [1, 4], "b": [10, 40]}) + >>> df1.intersect(df2).to_pydict() + {'a': [1], 'b': [10]} + + Intersect with deduplication: + + >>> df1 = ctx.from_pydict({"a": [1, 1, 2], "b": [10, 10, 20]}) + >>> df2 = ctx.from_pydict({"a": [1, 1], "b": [10, 10]}) + >>> df1.intersect(df2, distinct=True).to_pydict() + {'a': [1], 'b': [10]} """ - return DataFrame(self.df.union_distinct(other.df)) + return DataFrame(self.df.intersect(other.df, distinct)) - def intersect(self, other: DataFrame) -> DataFrame: - """Calculate the intersection of two :py:class:`DataFrame`. + def except_all(self, other: DataFrame, distinct: bool = False) -> DataFrame: + """Calculate the set difference of two :py:class:`DataFrame`. + + Returns rows that are in this DataFrame but not in ``other``. The two :py:class:`DataFrame` must have exactly the same schema. Args: - other: DataFrame to intersect with. + other: DataFrame to calculate exception with. + distinct: If ``True``, duplicate rows are removed from the result. Returns: - DataFrame after intersection. + DataFrame after set difference. + + Examples: + Remove rows present in ``df2``: + + >>> ctx = dfn.SessionContext() + >>> df1 = ctx.from_pydict({"a": [1, 2, 3], "b": [10, 20, 30]}) + >>> df2 = ctx.from_pydict({"a": [1, 2], "b": [10, 20]}) + >>> df1.except_all(df2).sort("a").to_pydict() + {'a': [3], 'b': [30]} + + Remove rows present in ``df2`` and deduplicate: + + >>> df1.except_all(df2, distinct=True).sort("a").to_pydict() + {'a': [3], 'b': [30]} """ - return DataFrame(self.df.intersect(other.df)) + return DataFrame(self.df.except_all(other.df, distinct)) - def except_all(self, other: DataFrame) -> DataFrame: - """Calculate the exception of two :py:class:`DataFrame`. + def union_by_name(self, other: DataFrame, distinct: bool = False) -> DataFrame: + """Union two :py:class:`DataFrame` matching columns by name. - The two :py:class:`DataFrame` must have exactly the same schema. + Unlike :py:meth:`union` which matches columns by position, this method + matches columns by their names, allowing DataFrames with different + column orders to be combined. Args: - other: DataFrame to calculate exception with. + other: DataFrame to union with. + distinct: If ``True``, duplicate rows are removed from the result. Returns: - DataFrame after exception. + DataFrame after union by name. + + Examples: + Combine DataFrames with different column orders: + + >>> ctx = dfn.SessionContext() + >>> df1 = ctx.from_pydict({"a": [1], "b": [10]}) + >>> df2 = ctx.from_pydict({"b": [20], "a": [2]}) + >>> df1.union_by_name(df2).sort("a").to_pydict() + {'a': [1, 2], 'b': [10, 20]} + + Union by name with deduplication: + + >>> df1 = ctx.from_pydict({"a": [1, 1], "b": [10, 10]}) + >>> df2 = ctx.from_pydict({"b": [10], "a": [1]}) + >>> df1.union_by_name(df2, distinct=True).to_pydict() + {'a': [1], 'b': [10]} """ - return DataFrame(self.df.except_all(other.df)) + return DataFrame(self.df.union_by_name(other.df, distinct)) + + def distinct_on( + self, + on_expr: list[Expr], + select_expr: list[Expr], + sort_expr: list[SortKey] | None = None, + ) -> DataFrame: + """Deduplicate rows based on specific columns. + + Returns a new DataFrame with one row per unique combination of the + ``on_expr`` columns, keeping the first row per group as determined by + ``sort_expr``. + + Args: + on_expr: Expressions that determine uniqueness. + select_expr: Expressions to include in the output. + sort_expr: Optional sort expressions to determine which row to keep. + + Returns: + DataFrame after deduplication. + + Examples: + Keep the row with the smallest ``b`` for each unique ``a``: + + >>> from datafusion import col + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 1, 2, 2], "b": [10, 20, 30, 40]}) + >>> df.distinct_on( + ... [col("a")], + ... [col("a"), col("b")], + ... [col("a").sort(ascending=True), col("b").sort(ascending=True)], + ... ).sort("a").to_pydict() + {'a': [1, 2], 'b': [10, 30]} + """ + on_raw = expr_list_to_raw_expr_list(on_expr) + select_raw = expr_list_to_raw_expr_list(select_expr) + sort_raw = sort_list_to_raw_sort_list(sort_expr) if sort_expr else None + return DataFrame(self.df.distinct_on(on_raw, select_raw, sort_raw)) + + def sort_by(self, *exprs: Expr | str) -> DataFrame: + """Sort the DataFrame by column expressions in ascending order. + + This is a convenience method that sorts the DataFrame by the given + expressions in ascending order with nulls last. For more control over + sort direction and null ordering, use :py:meth:`sort` instead. + + Args: + exprs: Expressions or column names to sort by. + + Returns: + DataFrame after sorting. + + Examples: + Sort by a single column: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [3, 1, 2]}) + >>> df.sort_by("a").to_pydict() + {'a': [1, 2, 3]} + """ + raw = [_to_raw_expr(e) for e in exprs] + return DataFrame(self.df.sort_by(raw)) def write_csv( self, @@ -1310,23 +1604,52 @@ def count(self) -> int: return self.df.count() @deprecated("Use :py:func:`unnest_columns` instead.") - def unnest_column(self, column: str, preserve_nulls: bool = True) -> DataFrame: + def unnest_column( + self, + column: str, + preserve_nulls: bool = True, + ) -> DataFrame: """See :py:func:`unnest_columns`.""" return DataFrame(self.df.unnest_column(column, preserve_nulls=preserve_nulls)) - def unnest_columns(self, *columns: str, preserve_nulls: bool = True) -> DataFrame: + def unnest_columns( + self, + *columns: str, + preserve_nulls: bool = True, + recursions: list[tuple[str, str, int]] | None = None, + ) -> DataFrame: """Expand columns of arrays into a single row per array element. Args: columns: Column names to perform unnest operation on. preserve_nulls: If False, rows with null entries will not be returned. + recursions: Optional list of ``(input_column, output_column, depth)`` + tuples that control how deeply nested columns are unnested. Any + column not mentioned here is unnested with depth 1. Returns: A DataFrame with the columns expanded. + + Examples: + Unnest an array column: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2], [3]], "b": ["x", "y"]}) + >>> df.unnest_columns("a").to_pydict() + {'a': [1, 2, 3], 'b': ['x', 'x', 'y']} + + With explicit recursion depth: + + >>> df.unnest_columns("a", recursions=[("a", "a", 1)]).to_pydict() + {'a': [1, 2, 3], 'b': ['x', 'x', 'y']} """ columns = list(columns) - return DataFrame(self.df.unnest_columns(columns, preserve_nulls=preserve_nulls)) + return DataFrame( + self.df.unnest_columns( + columns, preserve_nulls=preserve_nulls, recursions=recursions + ) + ) def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: """Export the DataFrame as an Arrow C Stream. diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index 759d6278c..bb8e9685c 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -29,6 +29,7 @@ import pytest from datafusion import ( DataFrame, + ExplainFormat, InsertOp, ParquetColumnOptions, ParquetWriterOptions, @@ -3569,3 +3570,263 @@ def test_read_parquet_file_sort_order(tmp_path, file_sort_order): pa.parquet.write_table(table, path) df = ctx.read_parquet(path, file_sort_order=file_sort_order) assert df.collect()[0].column(0).to_pylist() == [1, 2] + + +@pytest.mark.parametrize( + ("df1_data", "df2_data", "method", "kwargs", "expected_a", "expected_b"), + [ + pytest.param( + {"a": [1, 2, 3, 1], "b": [10, 20, 30, 10]}, + {"a": [1, 2], "b": [10, 20]}, + "except_all", + {"distinct": True}, + [3], + [30], + id="except_all(distinct=True): removes matching rows and deduplicates", + ), + pytest.param( + {"a": [1, 2, 3, 1], "b": [10, 20, 30, 10]}, + {"a": [1, 4], "b": [10, 40]}, + "intersect", + {"distinct": True}, + [1], + [10], + id="intersect(distinct=True): keeps common rows and deduplicates", + ), + pytest.param( + {"a": [1], "b": [10]}, + {"b": [20], "a": [2]}, # reversed column order tests matching by name + "union_by_name", + {}, + [1, 2], + [10, 20], + id="union_by_name: matches columns by name not position", + ), + ], +) +def test_set_operations_distinct( + df1_data, df2_data, method, kwargs, expected_a, expected_b +): + ctx = SessionContext() + df1 = ctx.from_pydict(df1_data) + df2 = ctx.from_pydict(df2_data) + result = ( + getattr(df1, method)(df2, **kwargs) + .sort(column("a").sort(ascending=True)) + .collect()[0] + ) + assert result.column(0).to_pylist() == expected_a + assert result.column(1).to_pylist() == expected_b + + +def test_union_by_name_distinct(): + ctx = SessionContext() + df1 = ctx.from_pydict({"a": [1, 1], "b": [10, 10]}) + df2 = ctx.from_pydict({"b": [10], "a": [1]}) + result = df1.union_by_name(df2, distinct=True).collect()[0] + assert result.column(0).to_pylist() == [1] + assert result.column(1).to_pylist() == [10] + + +def test_column_qualified(): + """DataFrame.column() returns a qualified column expression.""" + ctx = SessionContext() + df = ctx.from_pydict({"a": [1, 2], "b": [3, 4]}) + expr = df.column("a") + result = df.select(expr).collect()[0] + assert result.column(0).to_pylist() == [1, 2] + + +def test_column_not_found(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1]}) + with pytest.raises(Exception, match="not found"): + df.column("z") + + +def test_column_ambiguous(): + """After a join, duplicate column names that cannot be resolved raise an error.""" + ctx = SessionContext() + left = ctx.from_pydict({"id": [1, 2], "val": [10, 20]}) + right = ctx.from_pydict({"id": [1, 2], "val": [30, 40]}) + joined = left.join(right, on="id", how="inner") + with pytest.raises(Exception, match="not found"): + joined.column("val") + + +def test_column_after_join(): + """Qualified column works for non-ambiguous columns after a join.""" + ctx = SessionContext() + left = ctx.from_pydict({"id": [1, 2], "x": [10, 20]}) + right = ctx.from_pydict({"id": [1, 2], "y": [30, 40]}) + joined = left.join(right, on="id", how="inner") + expr = joined.column("y") + result = joined.select("id", expr).sort("id").collect()[0] + assert result.column(0).to_pylist() == [1, 2] + assert result.column(1).to_pylist() == [30, 40] + + +def test_col_join_disambiguate(): + """Use col() to disambiguate and select columns after a join.""" + ctx = SessionContext() + df1 = ctx.from_pydict({"foo": [1, 2, 3], "bar": [5, 6, 7]}) + df2 = ctx.from_pydict({"foo": [1, 2, 3], "baz": [8, 9, 10]}) + joined = df1.join_on(df2, df1.col("foo") == df2.col("foo"), how="inner") + result = ( + joined.select(df1.col("foo"), df1.col("bar"), df2.col("baz")) + .sort(df1.col("foo")) + .to_pydict() + ) + assert result["bar"] == [5, 6, 7] + assert result["baz"] == [8, 9, 10] + + +def test_find_qualified_columns(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1, 2], "b": [3, 4], "c": [5, 6]}) + exprs = df.find_qualified_columns("a", "c") + assert len(exprs) == 2 + result = df.select(*exprs).collect()[0] + assert result.column(0).to_pylist() == [1, 2] + assert result.column(1).to_pylist() == [5, 6] + + +def test_find_qualified_columns_not_found(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1]}) + with pytest.raises(Exception, match="not found"): + df.find_qualified_columns("a", "z") + + +def test_distinct_on(): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1, 1, 2, 2], "b": [10, 20, 30, 40]}) + result = ( + df.distinct_on( + [column("a")], + [column("a"), column("b")], + [column("a").sort(ascending=True), column("b").sort(ascending=True)], + ) + .sort(column("a").sort(ascending=True)) + .collect()[0] + ) + # Keeps the first row per group (smallest b per a) + assert result.column(0).to_pylist() == [1, 2] + assert result.column(1).to_pylist() == [10, 30] + + +@pytest.mark.parametrize( + ("input_values", "expected"), + [ + ([3, 1, 2], [1, 2, 3]), + ([1, 2, 3], [1, 2, 3]), + ([3, None, 1, 2], [1, 2, 3, None]), + ], +) +def test_sort_by(input_values, expected): + """sort_by always sorts ascending with nulls last regardless of input order.""" + ctx = SessionContext() + df = ctx.from_pydict({"a": input_values}) + result = df.sort_by(column("a")).collect()[0] + assert result.column(0).to_pylist() == expected + + +@pytest.mark.parametrize( + ("fmt", "verbose", "analyze", "expected_substring"), + [ + pytest.param(None, False, False, None, id="default format"), + pytest.param(ExplainFormat.TREE, False, False, "---", id="tree format"), + pytest.param( + ExplainFormat.INDENT, True, True, None, id="indent verbose+analyze" + ), + pytest.param(ExplainFormat.PGJSON, False, False, '"Plan"', id="pgjson format"), + pytest.param( + ExplainFormat.GRAPHVIZ, False, False, "digraph", id="graphviz format" + ), + ], +) +def test_explain_with_format(capsys, fmt, verbose, analyze, expected_substring): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1]}) + df.explain(verbose=verbose, analyze=analyze, format=fmt) + captured = capsys.readouterr() + assert "plan_type" in captured.out + if expected_substring is not None: + assert expected_substring in captured.out + + +@pytest.mark.parametrize( + ("window_exprs", "expected_columns"), + [ + pytest.param( + lambda: [ + f.row_number(partition_by=[column("b")], order_by=[column("a")]).alias( + "rn" + ), + ], + {"rn": [1, 2, 1]}, + id="single window expression", + ), + pytest.param( + lambda: [ + f.row_number(partition_by=[column("b")], order_by=[column("a")]).alias( + "rn" + ), + f.rank(partition_by=[column("b")], order_by=[column("a")]).alias("rnk"), + ], + {"rn": [1, 2, 1], "rnk": [1, 2, 1]}, + id="multiple window expressions", + ), + ], +) +def test_window(window_exprs, expected_columns): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1, 2, 3], "b": ["x", "x", "y"]}) + result = ( + df.window(*window_exprs()).sort(column("a").sort(ascending=True)).collect()[0] + ) + for col_name, expected_values in expected_columns.items(): + assert col_name in result.schema.names + assert ( + result.column(result.schema.get_field_index(col_name)).to_pylist() + == expected_values + ) + + +@pytest.mark.parametrize( + ("input_data", "recursions", "expected_a"), + [ + pytest.param( + {"a": [[1, 2], [3]], "b": ["x", "y"]}, + None, + [1, 2, 3], + id="basic unnest without recursions", + ), + pytest.param( + {"a": [[1, 2], [3]], "b": ["x", "y"]}, + [("a", "a", 1)], + [1, 2, 3], + id="explicit depth 1 matches basic unnest", + ), + pytest.param( + {"a": [[[1, 2], [3]], [[4]]], "b": ["x", "y"]}, + [("a", "a", 1)], + [[1, 2], [3], [4]], + id="depth 1 on nested lists keeps inner lists", + ), + pytest.param( + {"a": [[[1, 2], [3]], [[4]]], "b": ["x", "y"]}, + [("a", "a", 2)], + [1, 2, 3, 4], + id="depth 2 fully flattens nested lists", + ), + ], +) +def test_unnest_columns_with_recursions(input_data, recursions, expected_a): + ctx = SessionContext() + df = ctx.from_pydict(input_data) + kwargs = {} + if recursions is not None: + kwargs["recursions"] = recursions + result = df.unnest_columns("a", **kwargs).collect()[0] + assert result.column(0).to_pylist() == expected_a From 46f9ab8fcad03913234ce29e5075644c1ecdb9b7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 7 Apr 2026 15:03:38 -0400 Subject: [PATCH 14/83] Add missing deregister methods to SessionContext (#1473) * Add deregister methods to SessionContext for UDFs and object stores Expose upstream DataFusion deregister methods (deregister_udf, deregister_udaf, deregister_udwf, deregister_udtf, deregister_object_store) in both the Rust PyO3 bindings and Python wrappers, closing the gap identified in #1457. Co-Authored-By: Claude Opus 4.6 (1M context) * Fix deregister tests to expect ValueError instead of RuntimeError DataFusion raises ValueError for planning errors when a deregistered function is used in a query. Co-Authored-By: Claude Opus 4.6 (1M context) * Replace .unwrap() with proper error propagation in object store methods Url::parse() can fail on invalid input. Use .map_err() to convert the error into a Python exception instead of panicking. Co-Authored-By: Claude Opus 4.6 (1M context) * Minor move of import statement --------- Co-authored-by: Claude Opus 4.6 (1M context) --- crates/core/src/context.rs | 32 +++++++++- python/datafusion/context.py | 41 ++++++++++++ python/tests/test_context.py | 120 +++++++++++++++++++++++++++++++++++ 3 files changed, 192 insertions(+), 1 deletion(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 53994d2f5..1300a1595 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -434,11 +434,25 @@ impl PySessionContext { &upstream_host }; let url_string = format!("{scheme}{derived_host}"); - let url = Url::parse(&url_string).unwrap(); + let url = Url::parse(&url_string).map_err(|e| PyValueError::new_err(e.to_string()))?; self.ctx.runtime_env().register_object_store(&url, store); Ok(()) } + /// Deregister an object store with the given url + #[pyo3(signature = (scheme, host=None))] + pub fn deregister_object_store( + &self, + scheme: &str, + host: Option<&str>, + ) -> PyDataFusionResult<()> { + let host = host.unwrap_or(""); + let url_string = format!("{scheme}{host}"); + let url = Url::parse(&url_string).map_err(|e| PyDataFusionError::Common(e.to_string()))?; + self.ctx.runtime_env().deregister_object_store(&url)?; + Ok(()) + } + #[allow(clippy::too_many_arguments)] #[pyo3(signature = (name, path, table_partition_cols=vec![], file_extension=".parquet", @@ -492,6 +506,10 @@ impl PySessionContext { self.ctx.register_udtf(&name, func); } + pub fn deregister_udtf(&self, name: &str) { + self.ctx.deregister_udtf(name); + } + #[pyo3(signature = (query, options=None, param_values=HashMap::default(), param_strings=HashMap::default()))] pub fn sql_with_options( &self, @@ -975,16 +993,28 @@ impl PySessionContext { Ok(()) } + pub fn deregister_udf(&self, name: &str) { + self.ctx.deregister_udf(name); + } + pub fn register_udaf(&self, udaf: PyAggregateUDF) -> PyResult<()> { self.ctx.register_udaf(udaf.function); Ok(()) } + pub fn deregister_udaf(&self, name: &str) { + self.ctx.deregister_udaf(name); + } + pub fn register_udwf(&self, udwf: PyWindowUDF) -> PyResult<()> { self.ctx.register_udwf(udwf.function); Ok(()) } + pub fn deregister_udwf(&self, name: &str) { + self.ctx.deregister_udwf(name); + } + #[pyo3(signature = (name="datafusion"))] pub fn catalog(&self, py: Python, name: &str) -> PyResult> { let catalog = self.ctx.catalog(name).ok_or(PyKeyError::new_err(format!( diff --git a/python/datafusion/context.py b/python/datafusion/context.py index c8edc816f..f190e3ca1 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -568,6 +568,15 @@ def register_object_store( """ self.ctx.register_object_store(schema, store, host) + def deregister_object_store(self, schema: str, host: str | None = None) -> None: + """Remove an object store from the session. + + Args: + schema: The data source schema (e.g. ``"s3://"``). + host: URL for the host (e.g. bucket name). + """ + self.ctx.deregister_object_store(schema, host) + def register_listing_table( self, name: str, @@ -894,6 +903,14 @@ def register_udtf(self, func: TableFunction) -> None: """Register a user defined table function.""" self.ctx.register_udtf(func._udtf) + def deregister_udtf(self, name: str) -> None: + """Remove a user-defined table function from the session. + + Args: + name: Name of the UDTF to deregister. + """ + self.ctx.deregister_udtf(name) + def register_record_batches( self, name: str, partitions: list[list[pa.RecordBatch]] ) -> None: @@ -1105,14 +1122,38 @@ def register_udf(self, udf: ScalarUDF) -> None: """Register a user-defined function (UDF) with the context.""" self.ctx.register_udf(udf._udf) + def deregister_udf(self, name: str) -> None: + """Remove a user-defined scalar function from the session. + + Args: + name: Name of the UDF to deregister. + """ + self.ctx.deregister_udf(name) + def register_udaf(self, udaf: AggregateUDF) -> None: """Register a user-defined aggregation function (UDAF) with the context.""" self.ctx.register_udaf(udaf._udaf) + def deregister_udaf(self, name: str) -> None: + """Remove a user-defined aggregate function from the session. + + Args: + name: Name of the UDAF to deregister. + """ + self.ctx.deregister_udaf(name) + def register_udwf(self, udwf: WindowUDF) -> None: """Register a user-defined window function (UDWF) with the context.""" self.ctx.register_udwf(udwf._udwf) + def deregister_udwf(self, name: str) -> None: + """Remove a user-defined window function from the session. + + Args: + name: Name of the UDWF to deregister. + """ + self.ctx.deregister_udwf(name) + def catalog(self, name: str = "datafusion") -> Catalog: """Retrieve a catalog by name.""" return Catalog(self.ctx.catalog(name)) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 5df6ed20f..8491cc3a5 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -31,6 +31,7 @@ Table, column, literal, + udf, ) @@ -351,6 +352,125 @@ def test_deregister_table(ctx, database): assert public.names() == {"csv1", "csv2"} +def test_deregister_udf(): + ctx = SessionContext() + + is_null = udf( + lambda x: x.is_null(), + [pa.float64()], + pa.bool_(), + volatility="immutable", + name="my_is_null", + ) + ctx.register_udf(is_null) + + # Verify it works + df = ctx.from_pydict({"a": [1.0, None]}) + ctx.register_table("t", df.into_view()) + result = ctx.sql("SELECT my_is_null(a) FROM t").collect() + assert result[0].column(0) == pa.array([False, True]) + + # Deregister and verify it's gone + ctx.deregister_udf("my_is_null") + with pytest.raises(ValueError): + ctx.sql("SELECT my_is_null(a) FROM t").collect() + + +def test_deregister_udaf(): + import pyarrow.compute as pc + + ctx = SessionContext() + from datafusion import Accumulator, udaf + + class MySum(Accumulator): + def __init__(self): + self._sum = 0.0 + + def update(self, values: pa.Array) -> None: + self._sum += pc.sum(values).as_py() + + def merge(self, states: list[pa.Array]) -> None: + self._sum += pc.sum(states[0]).as_py() + + def state(self) -> list: + return [self._sum] + + def evaluate(self) -> pa.Scalar: + return self._sum + + my_sum = udaf( + MySum, + [pa.float64()], + pa.float64(), + [pa.float64()], + volatility="immutable", + name="my_sum", + ) + ctx.register_udaf(my_sum) + df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + ctx.register_table("t", df.into_view()) + + result = ctx.sql("SELECT my_sum(a) FROM t").collect() + assert result[0].column(0) == pa.array([6.0]) + + ctx.deregister_udaf("my_sum") + with pytest.raises(ValueError): + ctx.sql("SELECT my_sum(a) FROM t").collect() + + +def test_deregister_udwf(): + ctx = SessionContext() + from datafusion import udwf + from datafusion.user_defined import WindowEvaluator + + class MyRowNumber(WindowEvaluator): + def __init__(self): + self._row = 0 + + def evaluate_all(self, values, num_rows): + return pa.array(list(range(1, num_rows + 1)), type=pa.uint64()) + + my_row_number = udwf( + MyRowNumber, + [pa.float64()], + pa.uint64(), + volatility="immutable", + name="my_row_number", + ) + ctx.register_udwf(my_row_number) + df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + ctx.register_table("t", df.into_view()) + + result = ctx.sql("SELECT my_row_number(a) OVER () FROM t").collect() + assert result[0].column(0) == pa.array([1, 2, 3], type=pa.uint64()) + + ctx.deregister_udwf("my_row_number") + with pytest.raises(ValueError): + ctx.sql("SELECT my_row_number(a) OVER () FROM t").collect() + + +def test_deregister_udtf(): + import pyarrow.dataset as ds + + ctx = SessionContext() + from datafusion import Table, udtf + + class MyTable: + def __call__(self): + batch = pa.RecordBatch.from_pydict({"x": [1, 2, 3]}) + return Table(ds.dataset([batch])) + + my_table = udtf(MyTable(), "my_table") + ctx.register_udtf(my_table) + + result = ctx.sql("SELECT * FROM my_table()").collect() + assert result[0].column(0) == pa.array([1, 2, 3]) + + ctx.deregister_udtf("my_table") + with pytest.raises(ValueError): + ctx.sql("SELECT * FROM my_table()").collect() + + def test_register_table_from_dataframe(ctx): df = ctx.from_pydict({"a": [1, 2]}) ctx.register_table("df_tbl", df) From aa3b1948c3a49d14395093287a6e93354229c539 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 8 Apr 2026 09:22:28 -0400 Subject: [PATCH 15/83] Add missing registration methods (#1474) * Add missing SessionContext read/register methods for Arrow IPC and batches Add read_arrow, read_empty, register_arrow, and register_batch methods to SessionContext, exposing upstream DataFusion v53 functionality. The write_* methods and read_batch/read_batches are already covered by DataFrame.write_* and SessionContext.from_arrow respectively. Closes #1458. Co-Authored-By: Claude Opus 4.6 (1M context) * Remove redundant read_empty Rust binding, make Python read_empty an alias for empty_table Co-Authored-By: Claude Opus 4.6 (1M context) * Add pathlib.Path and empty batch tests for Arrow IPC and register_batch Co-Authored-By: Claude Opus 4.6 (1M context) * Make test_read_empty more robust with length and num_rows checks Co-Authored-By: Claude Opus 4.6 (1M context) * Add examples to docstrings for new register/read methods Co-Authored-By: Claude Opus 4.6 (1M context) * Empty table actually returns record batch of length one but there are no columns * Add optional argument examples to register_arrow and read_arrow docstrings Demonstrate schema= and file_extension= keyword arguments in the docstring examples for register_arrow and read_arrow, following project guidelines for optional parameter documentation. Co-Authored-By: Claude Opus 4.6 (1M context) * Simplify read_empty docstring to use alias pattern Follow the same See Also alias convention used in functions.py since read_empty is a simple alias for empty_table. Co-Authored-By: Claude Opus 4.6 (1M context) * Remove shared ctx from doctest namespace, use inline SessionContext Avoid shared SessionContext state across doctests by having each docstring example create its own ctx instance, matching the pattern used throughout the rest of the codebase. Co-Authored-By: Claude Opus 4.6 (1M context) * Remove redundant import pyarrow as pa from docstrings The pa alias is already provided by the doctest namespace in conftest.py, so inline imports are unnecessary. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- conftest.py | 2 + crates/core/src/context.rs | 58 +++++++++- python/datafusion/context.py | 181 ++++++++++++++++++++++++++++++ python/datafusion/user_defined.py | 3 - python/tests/test_context.py | 62 ++++++++++ 5 files changed, 302 insertions(+), 4 deletions(-) diff --git a/conftest.py b/conftest.py index 73e90077a..0c9410636 100644 --- a/conftest.py +++ b/conftest.py @@ -19,6 +19,7 @@ import datafusion as dfn import numpy as np +import pyarrow as pa import pytest from datafusion import col, lit from datafusion import functions as F @@ -29,6 +30,7 @@ def _doctest_namespace(doctest_namespace: dict) -> None: """Add common imports to the doctest namespace.""" doctest_namespace["dfn"] = dfn doctest_namespace["np"] = np + doctest_namespace["pa"] = pa doctest_namespace["col"] = col doctest_namespace["lit"] = lit doctest_namespace["F"] = F diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 1300a1595..ce11ef04e 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -41,7 +41,7 @@ use datafusion::execution::context::{ }; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, UnboundedMemoryPool}; -use datafusion::execution::options::ReadOptions; +use datafusion::execution::options::{ArrowReadOptions, ReadOptions}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::prelude::{ @@ -974,6 +974,39 @@ impl PySessionContext { Ok(()) } + #[pyo3(signature = (name, path, schema=None, file_extension=".arrow", table_partition_cols=vec![]))] + pub fn register_arrow( + &self, + name: &str, + path: &str, + schema: Option>, + file_extension: &str, + table_partition_cols: Vec<(String, PyArrowType)>, + py: Python, + ) -> PyDataFusionResult<()> { + let mut options = ArrowReadOptions::default().table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ); + options.file_extension = file_extension; + options.schema = schema.as_ref().map(|x| &x.0); + + let result = self.ctx.register_arrow(name, path, options); + wait_for_future(py, result)??; + Ok(()) + } + + pub fn register_batch( + &self, + name: &str, + batch: PyArrowType, + ) -> PyDataFusionResult<()> { + self.ctx.register_batch(name, batch.0)?; + Ok(()) + } + // Registers a PyArrow.Dataset pub fn register_dataset( &self, @@ -1214,6 +1247,29 @@ impl PySessionContext { Ok(PyDataFrame::new(df)) } + #[pyo3(signature = (path, schema=None, file_extension=".arrow", table_partition_cols=vec![]))] + pub fn read_arrow( + &self, + path: &str, + schema: Option>, + file_extension: &str, + table_partition_cols: Vec<(String, PyArrowType)>, + py: Python, + ) -> PyDataFusionResult { + let mut options = ArrowReadOptions::default().table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ); + options.file_extension = file_extension; + options.schema = schema.as_ref().map(|x| &x.0); + + let result = self.ctx.read_arrow(path, options); + let df = wait_for_future(py, result)??; + Ok(PyDataFrame::new(df)) + } + pub fn read_table(&self, table: Bound<'_, PyAny>) -> PyDataFusionResult { let session = self.clone().into_bound_py_any(table.py())?; let table = PyTable::new(table, Some(session))?; diff --git a/python/datafusion/context.py b/python/datafusion/context.py index f190e3ca1..7a306f04c 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -903,6 +903,27 @@ def register_udtf(self, func: TableFunction) -> None: """Register a user defined table function.""" self.ctx.register_udtf(func._udtf) + def register_batch(self, name: str, batch: pa.RecordBatch) -> None: + """Register a single :py:class:`pa.RecordBatch` as a table. + + Args: + name: Name of the resultant table. + batch: Record batch to register as a table. + + Examples: + >>> ctx = dfn.SessionContext() + >>> batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3]}) + >>> ctx.register_batch("batch_tbl", batch) + >>> ctx.sql("SELECT * FROM batch_tbl").collect()[0].column(0) + + [ + 1, + 2, + 3 + ] + """ + self.ctx.register_batch(name, batch) + def deregister_udtf(self, name: str) -> None: """Remove a user-defined table function from the session. @@ -1109,6 +1130,86 @@ def register_avro( name, str(path), schema, file_extension, table_partition_cols ) + def register_arrow( + self, + name: str, + path: str | pathlib.Path, + schema: pa.Schema | None = None, + file_extension: str = ".arrow", + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, + ) -> None: + """Register an Arrow IPC file as a table. + + The registered table can be referenced from SQL statements executed + against this context. + + Args: + name: Name of the table to register. + path: Path to the Arrow IPC file. + schema: The data source schema. + file_extension: File extension to select. + table_partition_cols: Partition columns. + + Examples: + >>> import tempfile, os + >>> ctx = dfn.SessionContext() + >>> table = pa.table({"x": [10, 20, 30]}) + >>> with tempfile.TemporaryDirectory() as tmpdir: + ... path = os.path.join(tmpdir, "data.arrow") + ... with pa.ipc.new_file(path, table.schema) as writer: + ... writer.write_table(table) + ... ctx.register_arrow("arrow_tbl", path) + ... ctx.sql("SELECT * FROM arrow_tbl").collect()[0].column(0) + + [ + 10, + 20, + 30 + ] + + Provide an explicit ``schema`` to override schema inference: + + >>> with tempfile.TemporaryDirectory() as tmpdir: + ... path = os.path.join(tmpdir, "data.arrow") + ... with pa.ipc.new_file(path, table.schema) as writer: + ... writer.write_table(table) + ... ctx.register_arrow( + ... "arrow_schema", + ... path, + ... schema=pa.schema([("x", pa.int64())]), + ... ) + ... ctx.sql("SELECT * FROM arrow_schema").collect()[0].column(0) + + [ + 10, + 20, + 30 + ] + + Use ``file_extension`` to read files with a non-default extension: + + >>> with tempfile.TemporaryDirectory() as tmpdir: + ... path = os.path.join(tmpdir, "data.ipc") + ... with pa.ipc.new_file(path, table.schema) as writer: + ... writer.write_table(table) + ... ctx.register_arrow( + ... "arrow_ipc", path, file_extension=".ipc" + ... ) + ... ctx.sql("SELECT * FROM arrow_ipc").collect()[0].column(0) + + [ + 10, + 20, + 30 + ] + """ + if table_partition_cols is None: + table_partition_cols = [] + table_partition_cols = _convert_table_partition_cols(table_partition_cols) + self.ctx.register_arrow( + name, str(path), schema, file_extension, table_partition_cols + ) + def register_dataset(self, name: str, dataset: pa.dataset.Dataset) -> None: """Register a :py:class:`pa.dataset.Dataset` as a table. @@ -1369,6 +1470,86 @@ def read_avro( self.ctx.read_avro(str(path), schema, file_partition_cols, file_extension) ) + def read_arrow( + self, + path: str | pathlib.Path, + schema: pa.Schema | None = None, + file_extension: str = ".arrow", + file_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, + ) -> DataFrame: + """Create a :py:class:`DataFrame` for reading an Arrow IPC data source. + + Args: + path: Path to the Arrow IPC file. + schema: The data source schema. + file_extension: File extension to select. + file_partition_cols: Partition columns. + + Returns: + DataFrame representation of the read Arrow IPC file. + + Examples: + >>> import tempfile, os + >>> ctx = dfn.SessionContext() + >>> table = pa.table({"a": [1, 2, 3]}) + >>> with tempfile.TemporaryDirectory() as tmpdir: + ... path = os.path.join(tmpdir, "data.arrow") + ... with pa.ipc.new_file(path, table.schema) as writer: + ... writer.write_table(table) + ... df = ctx.read_arrow(path) + ... df.collect()[0].column(0) + + [ + 1, + 2, + 3 + ] + + Provide an explicit ``schema`` to override schema inference: + + >>> with tempfile.TemporaryDirectory() as tmpdir: + ... path = os.path.join(tmpdir, "data.arrow") + ... with pa.ipc.new_file(path, table.schema) as writer: + ... writer.write_table(table) + ... df = ctx.read_arrow(path, schema=pa.schema([("a", pa.int64())])) + ... df.collect()[0].column(0) + + [ + 1, + 2, + 3 + ] + + Use ``file_extension`` to read files with a non-default extension: + + >>> with tempfile.TemporaryDirectory() as tmpdir: + ... path = os.path.join(tmpdir, "data.ipc") + ... with pa.ipc.new_file(path, table.schema) as writer: + ... writer.write_table(table) + ... df = ctx.read_arrow(path, file_extension=".ipc") + ... df.collect()[0].column(0) + + [ + 1, + 2, + 3 + ] + """ + if file_partition_cols is None: + file_partition_cols = [] + file_partition_cols = _convert_table_partition_cols(file_partition_cols) + return DataFrame( + self.ctx.read_arrow(str(path), schema, file_extension, file_partition_cols) + ) + + def read_empty(self) -> DataFrame: + """Create an empty :py:class:`DataFrame` with no columns or rows. + + See Also: + This is an alias for :meth:`empty_table`. + """ + return self.empty_table() + def read_table( self, table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset ) -> DataFrame: diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index 3eaccdfa3..848ab4cee 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -213,7 +213,6 @@ def udf(*args: Any, **kwargs: Any): # noqa: D417 Examples: Using ``udf`` as a function: - >>> import pyarrow as pa >>> import pyarrow.compute as pc >>> from datafusion.user_defined import ScalarUDF >>> def double_func(x): @@ -480,7 +479,6 @@ def udaf(*args: Any, **kwargs: Any): # noqa: D417, C901 instance in which this UDAF is used. Examples: - >>> import pyarrow as pa >>> import pyarrow.compute as pc >>> from datafusion.user_defined import AggregateUDF, Accumulator, udaf >>> class Summarize(Accumulator): @@ -874,7 +872,6 @@ def udwf(*args: Any, **kwargs: Any): # noqa: D417 When using ``udwf`` as a decorator, do not pass ``func`` explicitly. Examples: - >>> import pyarrow as pa >>> from datafusion.user_defined import WindowUDF, WindowEvaluator, udwf >>> class BiasedNumbers(WindowEvaluator): ... def __init__(self, start: int = 0): diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 8491cc3a5..25f66a647 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -788,6 +788,68 @@ def test_read_avro(ctx): assert avro_df is not None +def test_read_arrow(ctx, tmp_path): + # Write an Arrow IPC file, then read it back + table = pa.table({"a": [1, 2, 3], "b": ["x", "y", "z"]}) + arrow_path = tmp_path / "test.arrow" + with pa.ipc.new_file(str(arrow_path), table.schema) as writer: + writer.write_table(table) + + df = ctx.read_arrow(str(arrow_path)) + result = df.collect() + assert result[0].column(0) == pa.array([1, 2, 3]) + assert result[0].column(1) == pa.array(["x", "y", "z"]) + + # Also verify pathlib.Path works + df = ctx.read_arrow(arrow_path) + result = df.collect() + assert result[0].column(0) == pa.array([1, 2, 3]) + + +def test_read_empty(ctx): + df = ctx.read_empty() + result = df.collect() + assert len(result) == 1 + assert result[0].num_columns == 0 + + df = ctx.empty_table() + result = df.collect() + assert len(result) == 1 + assert result[0].num_columns == 0 + + +def test_register_arrow(ctx, tmp_path): + # Write an Arrow IPC file, then register and query it + table = pa.table({"x": [10, 20, 30]}) + arrow_path = tmp_path / "test.arrow" + with pa.ipc.new_file(str(arrow_path), table.schema) as writer: + writer.write_table(table) + + ctx.register_arrow("arrow_tbl", str(arrow_path)) + result = ctx.sql("SELECT * FROM arrow_tbl").collect() + assert result[0].column(0) == pa.array([10, 20, 30]) + + # Also verify pathlib.Path works + ctx.register_arrow("arrow_tbl_path", arrow_path) + result = ctx.sql("SELECT * FROM arrow_tbl_path").collect() + assert result[0].column(0) == pa.array([10, 20, 30]) + + +def test_register_batch(ctx): + batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}) + ctx.register_batch("batch_tbl", batch) + result = ctx.sql("SELECT * FROM batch_tbl").collect() + assert result[0].column(0) == pa.array([1, 2, 3]) + assert result[0].column(1) == pa.array([4, 5, 6]) + + +def test_register_batch_empty(ctx): + batch = pa.RecordBatch.from_pydict({"a": pa.array([], type=pa.int64())}) + ctx.register_batch("empty_batch_tbl", batch) + result = ctx.sql("SELECT * FROM empty_batch_tbl").collect() + assert result[0].num_rows == 0 + + def test_create_sql_options(): SQLOptions() From ecd14c10aff67169f2bfe1b7f86ff07621088dd0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 8 Apr 2026 11:11:48 -0400 Subject: [PATCH 16/83] Add missing SessionContext utility methods (#1475) * Add missing SessionContext utility methods Expose upstream DataFusion v53 utility methods: session_start_time, enable_ident_normalization, parse_sql_expr, execute_logical_plan, refresh_catalogs, remove_optimizer_rule, and table_provider. The add_optimizer_rule and add_analyzer_rule methods are omitted as the OptimizerRule and AnalyzerRule traits are not yet exposed to Python. Closes #1459. Co-Authored-By: Claude Opus 4.6 (1M context) * Raise KeyError from table_provider for consistency with table() Co-Authored-By: Claude Opus 4.6 (1M context) * Add docstring examples for new SessionContext utility methods Co-Authored-By: Claude Opus 4.6 (1M context) * update docstring * Address PR review feedback for SessionContext utility methods - Improve docstring examples to show actual output instead of asserts - Use doctest +SKIP for non-deterministic session_start_time output - Fix table_provider error mapping: outer async error is now RuntimeError - Strengthen tests: validate RFC 3339 with fromisoformat, test both optimizer rule removal paths, exact string match for parse_sql_expr, verify enable_ident_normalization with dynamic state change Co-Authored-By: Claude Opus 4.6 (1M context) * Fix test_session_start_time failure on Python 3.10 datetime.fromisoformat() only supports up to 6 fractional-second digits (microseconds) on Python 3.10. Truncate nanosecond precision before parsing. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- crates/core/src/context.rs | 49 ++++++++++++++- python/datafusion/context.py | 118 ++++++++++++++++++++++++++++++++++- python/tests/test_context.py | 55 ++++++++++++++++ 3 files changed, 219 insertions(+), 3 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index ce11ef04e..b4fe524df 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -28,7 +28,7 @@ use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory}; -use datafusion::common::{ScalarValue, TableReference, exec_err}; +use datafusion::common::{DFSchema, ScalarValue, TableReference, exec_err}; use datafusion::datasource::file_format::file_compression_type::FileCompressionType; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ @@ -60,7 +60,7 @@ use datafusion_python_util::{ }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; -use pyo3::exceptions::{PyKeyError, PyValueError}; +use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; use url::Url; @@ -70,11 +70,13 @@ use crate::catalog::{ PyCatalog, PyCatalogList, RustWrappedPyCatalogProvider, RustWrappedPyCatalogProviderList, }; use crate::common::data_type::PyScalarValue; +use crate::common::df_schema::PyDFSchema; use crate::dataframe::PyDataFrame; use crate::dataset::Dataset; use crate::errors::{ PyDataFusionError, PyDataFusionResult, from_datafusion_error, py_datafusion_err, }; +use crate::expr::PyExpr; use crate::expr::sort_expr::PySortExpr; use crate::options::PyCsvReadOptions; use crate::physical_plan::PyExecutionPlan; @@ -1113,6 +1115,49 @@ impl PySessionContext { self.ctx.session_id() } + pub fn session_start_time(&self) -> String { + self.ctx.session_start_time().to_rfc3339() + } + + pub fn enable_ident_normalization(&self) -> bool { + self.ctx.enable_ident_normalization() + } + + pub fn parse_sql_expr(&self, sql: &str, schema: PyDFSchema) -> PyDataFusionResult { + let df_schema: DFSchema = schema.into(); + Ok(self.ctx.parse_sql_expr(sql, &df_schema)?.into()) + } + + pub fn execute_logical_plan( + &self, + plan: PyLogicalPlan, + py: Python, + ) -> PyDataFusionResult { + let df = wait_for_future( + py, + self.ctx.execute_logical_plan(plan.plan.as_ref().clone()), + )??; + Ok(PyDataFrame::new(df)) + } + + pub fn refresh_catalogs(&self, py: Python) -> PyDataFusionResult<()> { + wait_for_future(py, self.ctx.refresh_catalogs())??; + Ok(()) + } + + pub fn remove_optimizer_rule(&self, name: &str) -> bool { + self.ctx.remove_optimizer_rule(name) + } + + pub fn table_provider(&self, name: &str, py: Python) -> PyResult { + let provider = wait_for_future(py, self.ctx.table_provider(name)) + // Outer error: runtime/async failure + .map_err(|e| PyRuntimeError::new_err(e.to_string()))? + // Inner error: table not found + .map_err(|e| PyKeyError::new_err(e.to_string()))?; + Ok(PyTable { table: provider }) + } + #[allow(clippy::too_many_arguments)] #[pyo3(signature = (path, schema=None, schema_infer_max_records=1000, file_extension=".json", table_partition_cols=vec![], file_compression_type=None))] pub fn read_json( diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 7a306f04c..e3949de83 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -63,7 +63,8 @@ import polars as pl # type: ignore[import] from datafusion.catalog import CatalogProvider, Table - from datafusion.expr import SortKey + from datafusion.common import DFSchema + from datafusion.expr import Expr, SortKey from datafusion.plan import ExecutionPlan, LogicalPlan from datafusion.user_defined import ( AggregateUDF, @@ -1283,6 +1284,121 @@ def session_id(self) -> str: """Return an id that uniquely identifies this :py:class:`SessionContext`.""" return self.ctx.session_id() + def session_start_time(self) -> str: + """Return the session start time as an RFC 3339 formatted string. + + Examples: + >>> ctx = SessionContext() + >>> ctx.session_start_time() # doctest: +SKIP + '2026-01-01T12:34:56.123456789+00:00' + """ + return self.ctx.session_start_time() + + def enable_ident_normalization(self) -> bool: + """Return whether identifier normalization (lowercasing) is enabled. + + Examples: + >>> ctx = SessionContext() + >>> ctx.enable_ident_normalization() + True + """ + return self.ctx.enable_ident_normalization() + + def parse_sql_expr(self, sql: str, schema: DFSchema) -> Expr: + """Parse a SQL expression string into a logical expression. + + Args: + sql: SQL expression string. + schema: Schema to use for resolving column references. + + Returns: + Parsed expression. + + Examples: + >>> from datafusion.common import DFSchema + >>> ctx = SessionContext() + >>> schema = DFSchema.empty() + >>> ctx.parse_sql_expr("1 + 2", schema=schema) + Expr(Int64(1) + Int64(2)) + """ + from datafusion.expr import Expr # noqa: PLC0415 + + return Expr(self.ctx.parse_sql_expr(sql, schema)) + + def execute_logical_plan(self, plan: LogicalPlan) -> DataFrame: + """Execute a :py:class:`~datafusion.plan.LogicalPlan` and return a DataFrame. + + Args: + plan: Logical plan to execute. + + Returns: + DataFrame resulting from the execution. + + Examples: + >>> ctx = SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> plan = df.logical_plan() + >>> df2 = ctx.execute_logical_plan(plan) + >>> df2.collect()[0].column(0) + + [ + 1, + 2, + 3 + ] + """ + return DataFrame(self.ctx.execute_logical_plan(plan._raw_plan)) + + def refresh_catalogs(self) -> None: + """Refresh catalog metadata. + + Examples: + >>> ctx = SessionContext() + >>> ctx.refresh_catalogs() + """ + self.ctx.refresh_catalogs() + + def remove_optimizer_rule(self, name: str) -> bool: + """Remove an optimizer rule by name. + + Args: + name: Name of the optimizer rule to remove. + + Returns: + True if a rule with the given name was found and removed. + + Examples: + >>> ctx = SessionContext() + >>> ctx.remove_optimizer_rule("nonexistent_rule") + False + """ + return self.ctx.remove_optimizer_rule(name) + + def table_provider(self, name: str) -> Table: + """Return the :py:class:`~datafusion.catalog.Table` for the given table name. + + Args: + name: Name of the table. + + Returns: + The table provider. + + Raises: + KeyError: If the table is not found. + + Examples: + >>> import pyarrow as pa + >>> ctx = SessionContext() + >>> batch = pa.RecordBatch.from_pydict({"x": [1, 2]}) + >>> ctx.register_record_batches("my_table", [[batch]]) + >>> tbl = ctx.table_provider("my_table") + >>> tbl.schema + x: int64 + """ + from datafusion.catalog import Table # noqa: PLC0415 + + return Table(self.ctx.table_provider(name)) + def read_json( self, path: str | pathlib.Path, diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 25f66a647..13c05a9e6 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -671,6 +671,61 @@ def test_table_not_found(ctx): ctx.table(f"not-found-{uuid4()}") +def test_session_start_time(ctx): + import datetime + import re + + st = ctx.session_start_time() + assert isinstance(st, str) + # Truncate nanoseconds to microseconds for Python 3.10 compat + st = re.sub(r"(\.\d{6})\d+", r"\1", st) + dt = datetime.datetime.fromisoformat(st) + assert dt.isoformat() + + +def test_enable_ident_normalization(ctx): + assert ctx.enable_ident_normalization() is True + ctx.sql("SET datafusion.sql_parser.enable_ident_normalization = false") + assert ctx.enable_ident_normalization() is False + + +def test_parse_sql_expr(ctx): + from datafusion.common import DFSchema + + schema = DFSchema.empty() + expr = ctx.parse_sql_expr("1 + 2", schema) + assert str(expr) == "Expr(Int64(1) + Int64(2))" + + +def test_execute_logical_plan(ctx): + df = ctx.from_pydict({"a": [1, 2, 3]}) + plan = df.logical_plan() + df2 = ctx.execute_logical_plan(plan) + result = df2.collect() + assert result[0].column(0) == pa.array([1, 2, 3]) + + +def test_refresh_catalogs(ctx): + ctx.refresh_catalogs() + + +def test_remove_optimizer_rule(ctx): + assert ctx.remove_optimizer_rule("push_down_filter") is True + assert ctx.remove_optimizer_rule("nonexistent_rule") is False + + +def test_table_provider(ctx): + batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) + ctx.register_record_batches("provider_test", [[batch]]) + tbl = ctx.table_provider("provider_test") + assert tbl.schema == pa.schema([("x", pa.int64())]) + + +def test_table_provider_not_found(ctx): + with pytest.raises(KeyError): + ctx.table_provider("nonexistent_table") + + def test_read_json(ctx): path = pathlib.Path(__file__).parent.resolve() From 3585c11eed778810e3317c56c2c25a8cdc29be5b Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 9 Apr 2026 07:38:59 -0400 Subject: [PATCH 17/83] minor: remove deprecated interfaces (#1481) * udf module has been deprecated since DF47. html_formatter module has been deprecated since DF48. * database has been deprecated since DF48 * select_columns has been deprecated since DF43 * unnest_column has been deprecated since DF42 * display_name has been deprecated since DF42 * window() has been deprecated since DF50 * serde functions have been deprecated since DF42 * from_arrow_table and tables have been deprecated since DF42 * RuntimeConfig has been deprecated since DF44 * Update user documentation to remove deprecated function * update tpch examples for latest function uses * Remove unnecessary options in example * update rendering for the most recent dataframe_formatter instead of the deprecated html_formatter --- crates/core/src/context.rs | 15 -- crates/core/src/dataframe.rs | 29 +-- crates/core/src/functions.rs | 133 +---------- .../user-guide/common-operations/windows.rst | 15 +- .../source/user-guide/dataframe/rendering.rst | 225 ++++++++++-------- examples/tpch/q02_minimum_cost_supplier.py | 8 +- .../q11_important_stock_identification.py | 3 +- examples/tpch/q15_top_supplier.py | 4 +- examples/tpch/q17_small_quantity_order.py | 8 +- examples/tpch/q22_global_sales_opportunity.py | 4 +- python/datafusion/__init__.py | 3 +- python/datafusion/catalog.py | 10 - python/datafusion/context.py | 21 -- python/datafusion/dataframe.py | 20 -- python/datafusion/dataframe_formatter.py | 8 +- python/datafusion/expr.py | 15 -- python/datafusion/functions.py | 55 +---- python/datafusion/html_formatter.py | 29 --- python/datafusion/substrait.py | 25 -- python/datafusion/udf.py | 29 --- python/tests/test_expr.py | 21 -- 21 files changed, 154 insertions(+), 526 deletions(-) delete mode 100644 python/datafusion/html_formatter.py delete mode 100644 python/datafusion/udf.py diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index b4fe524df..e46d359d6 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1072,21 +1072,6 @@ impl PySessionContext { self.ctx.catalog_names().into_iter().collect() } - pub fn tables(&self) -> HashSet { - self.ctx - .catalog_names() - .into_iter() - .filter_map(|name| self.ctx.catalog(&name)) - .flat_map(move |catalog| { - catalog - .schema_names() - .into_iter() - .filter_map(move |name| catalog.schema(&name)) - }) - .flat_map(|schema| schema.table_names()) - .collect() - } - pub fn table(&self, name: &str, py: Python) -> PyResult { let res = wait_for_future(py, self.ctx.table(name)) .map_err(|e| PyKeyError::new_err(e.to_string()))?; diff --git a/crates/core/src/dataframe.rs b/crates/core/src/dataframe.rs index fff5118d5..c067eac30 100644 --- a/crates/core/src/dataframe.rs +++ b/crates/core/src/dataframe.rs @@ -468,17 +468,17 @@ impl PyDataFrame { fn __getitem__(&self, key: Bound<'_, PyAny>) -> PyDataFusionResult { if let Ok(key) = key.extract::() { // df[col] - self.select_columns(vec![key]) + self.select_exprs(vec![key]) } else if let Ok(tuple) = key.cast::() { // df[col1, col2, col3] let keys = tuple .iter() .map(|item| item.extract::()) .collect::>>()?; - self.select_columns(keys) + self.select_exprs(keys) } else if let Ok(keys) = key.extract::>() { // df[[col1, col2, col3]] - self.select_columns(keys) + self.select_exprs(keys) } else { let message = "DataFrame can only be indexed by string index or indices".to_string(); Err(PyDataFusionError::Common(message)) @@ -554,13 +554,6 @@ impl PyDataFrame { Ok(PyTable::from(table_provider)) } - #[pyo3(signature = (*args))] - fn select_columns(&self, args: Vec) -> PyDataFusionResult { - let args = args.iter().map(|s| s.as_ref()).collect::>(); - let df = self.df.as_ref().clone().select_columns(&args)?; - Ok(Self::new(df)) - } - #[pyo3(signature = (*args))] fn select_exprs(&self, args: Vec) -> PyDataFusionResult { let args = args.iter().map(|s| s.as_ref()).collect::>(); @@ -890,22 +883,6 @@ impl PyDataFrame { Ok(Self::new(new_df)) } - #[pyo3(signature = (column, preserve_nulls=true, recursions=None))] - fn unnest_column( - &self, - column: &str, - preserve_nulls: bool, - recursions: Option>, - ) -> PyDataFusionResult { - let unnest_options = build_unnest_options(preserve_nulls, recursions); - let df = self - .df - .as_ref() - .clone() - .unnest_columns_with_options(&[column], unnest_options)?; - Ok(Self::new(df)) - } - #[pyo3(signature = (columns, preserve_nulls=true, recursions=None))] fn unnest_columns( &self, diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index f173aaa51..7feb62d79 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -18,20 +18,14 @@ use std::collections::HashMap; use datafusion::common::{Column, ScalarValue, TableReference}; -use datafusion::execution::FunctionRegistry; -use datafusion::functions_aggregate::all_default_aggregate_functions; -use datafusion::functions_window::all_default_window_functions; -use datafusion::logical_expr::expr::{ - Alias, FieldMetadata, NullTreatment as DFNullTreatment, WindowFunction, WindowFunctionParams, -}; -use datafusion::logical_expr::{Expr, ExprFunctionExt, WindowFrame, WindowFunctionDefinition, lit}; +use datafusion::logical_expr::expr::{Alias, FieldMetadata, NullTreatment as DFNullTreatment}; +use datafusion::logical_expr::{Expr, ExprFunctionExt, lit}; use datafusion::{functions, functions_aggregate, functions_window}; use pyo3::prelude::*; use pyo3::wrap_pyfunction; use crate::common::data_type::{NullTreatment, PyScalarValue}; -use crate::context::PySessionContext; -use crate::errors::{PyDataFusionError, PyDataFusionResult}; +use crate::errors::PyDataFusionResult; use crate::expr::PyExpr; use crate::expr::conditional_expr::PyCaseBuilder; use crate::expr::sort_expr::{PySortExpr, to_sort_expressions}; @@ -306,126 +300,6 @@ fn when(when: PyExpr, then: PyExpr) -> PyResult { Ok(PyCaseBuilder::new(None).when(when, then)) } -/// Helper function to find the appropriate window function. -/// -/// Search procedure: -/// 1) Search built in window functions, which are being deprecated. -/// 1) If a session context is provided: -/// 1) search User Defined Aggregate Functions (UDAFs) -/// 1) search registered window functions -/// 1) search registered aggregate functions -/// 1) If no function has been found, search default aggregate functions. -/// -/// NOTE: we search the built-ins first because the `UDAF` versions currently do not have the same behavior. -fn find_window_fn( - name: &str, - ctx: Option, -) -> PyDataFusionResult { - if let Some(ctx) = ctx { - // search UDAFs - let udaf = ctx - .ctx - .udaf(name) - .map(WindowFunctionDefinition::AggregateUDF) - .ok(); - - if let Some(udaf) = udaf { - return Ok(udaf); - } - - let session_state = ctx.ctx.state(); - - // search registered window functions - let window_fn = session_state - .window_functions() - .get(name) - .map(|f| WindowFunctionDefinition::WindowUDF(f.clone())); - - if let Some(window_fn) = window_fn { - return Ok(window_fn); - } - - // search registered aggregate functions - let agg_fn = session_state - .aggregate_functions() - .get(name) - .map(|f| WindowFunctionDefinition::AggregateUDF(f.clone())); - - if let Some(agg_fn) = agg_fn { - return Ok(agg_fn); - } - } - - // search default aggregate functions - let agg_fn = all_default_aggregate_functions() - .iter() - .find(|v| v.name() == name || v.aliases().contains(&name.to_string())) - .map(|f| WindowFunctionDefinition::AggregateUDF(f.clone())); - - if let Some(agg_fn) = agg_fn { - return Ok(agg_fn); - } - - // search default window functions - let window_fn = all_default_window_functions() - .iter() - .find(|v| v.name() == name || v.aliases().contains(&name.to_string())) - .map(|f| WindowFunctionDefinition::WindowUDF(f.clone())); - - if let Some(window_fn) = window_fn { - return Ok(window_fn); - } - - Err(PyDataFusionError::Common(format!( - "window function `{name}` not found" - ))) -} - -/// Creates a new Window function expression -#[allow(clippy::too_many_arguments)] -#[pyfunction] -#[pyo3(signature = (name, args, partition_by=None, order_by=None, window_frame=None, filter=None, distinct=false, ctx=None))] -fn window( - name: &str, - args: Vec, - partition_by: Option>, - order_by: Option>, - window_frame: Option, - filter: Option, - distinct: bool, - ctx: Option, -) -> PyResult { - let fun = find_window_fn(name, ctx)?; - - let window_frame = window_frame - .map(|w| w.into()) - .unwrap_or(WindowFrame::new(order_by.as_ref().map(|v| !v.is_empty()))); - let filter = filter.map(|f| f.expr.into()); - - Ok(PyExpr { - expr: datafusion::logical_expr::Expr::WindowFunction(Box::new(WindowFunction { - fun, - params: WindowFunctionParams { - args: args.into_iter().map(|x| x.expr).collect::>(), - partition_by: partition_by - .unwrap_or_default() - .into_iter() - .map(|x| x.expr) - .collect::>(), - order_by: order_by - .unwrap_or_default() - .into_iter() - .map(|x| x.into()) - .collect::>(), - window_frame, - filter, - distinct, - null_treatment: None, - }, - })), - }) -} - // Generates a [pyo3] wrapper for associated aggregate functions. // All of the builder options are exposed to the python internal // function and we rely on the wrappers to only use those that @@ -1186,7 +1060,6 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(self::uuid))?; // Use self to avoid name collision m.add_wrapped(wrap_pyfunction!(var_pop))?; m.add_wrapped(wrap_pyfunction!(var_sample))?; - m.add_wrapped(wrap_pyfunction!(window))?; m.add_wrapped(wrap_pyfunction!(regr_avgx))?; m.add_wrapped(wrap_pyfunction!(regr_avgy))?; m.add_wrapped(wrap_pyfunction!(regr_count))?; diff --git a/docs/source/user-guide/common-operations/windows.rst b/docs/source/user-guide/common-operations/windows.rst index c8fdea8f4..d77881bcf 100644 --- a/docs/source/user-guide/common-operations/windows.rst +++ b/docs/source/user-guide/common-operations/windows.rst @@ -175,10 +175,7 @@ it's ``Type 2`` column that are null. Aggregate Functions ------------------- -You can use any :ref:`Aggregation Function` as a window function. Currently -aggregate functions must use the deprecated -:py:func:`datafusion.functions.window` API but this should be resolved in -DataFusion 42.0 (`Issue Link `_). Here +You can use any :ref:`Aggregation Function` as a window function. Here is an example that shows how to compare each pokemons’s attack power with the average attack power in its ``"Type 1"`` using the :py:func:`datafusion.functions.avg` function. @@ -189,10 +186,12 @@ power in its ``"Type 1"`` using the :py:func:`datafusion.functions.avg` function col('"Name"'), col('"Attack"'), col('"Type 1"'), - f.window("avg", [col('"Attack"')]) - .partition_by(col('"Type 1"')) - .build() - .alias("Average Attack"), + f.avg(col('"Attack"')).over( + Window( + window_frame=WindowFrame("rows", None, None), + partition_by=[col('"Type 1"')], + ) + ).alias("Average Attack"), ) Available Functions diff --git a/docs/source/user-guide/dataframe/rendering.rst b/docs/source/user-guide/dataframe/rendering.rst index 9dea948bb..dc61a422f 100644 --- a/docs/source/user-guide/dataframe/rendering.rst +++ b/docs/source/user-guide/dataframe/rendering.rst @@ -15,18 +15,18 @@ .. specific language governing permissions and limitations .. under the License. -HTML Rendering in Jupyter -========================= +DataFrame Rendering +=================== -When working in Jupyter notebooks or other environments that support rich HTML display, -DataFusion DataFrames automatically render as nicely formatted HTML tables. This functionality -is provided by the ``_repr_html_`` method, which is automatically called by Jupyter to provide -a richer visualization than plain text output. +DataFusion provides configurable rendering for DataFrames in both plain text and HTML +formats. The ``datafusion.dataframe_formatter`` module controls how DataFrames are +displayed in Jupyter notebooks (via ``_repr_html_``), in the terminal (via ``__repr__``), +and anywhere else a string or HTML representation is needed. -Basic HTML Rendering --------------------- +Basic Rendering +--------------- -In a Jupyter environment, simply displaying a DataFrame object will trigger HTML rendering: +In a Jupyter environment, displaying a DataFrame triggers HTML rendering: .. code-block:: python @@ -36,74 +36,117 @@ In a Jupyter environment, simply displaying a DataFrame object will trigger HTML # Explicit display also uses HTML rendering display(df) -Customizing HTML Rendering ---------------------------- +In a terminal or when converting to string, plain text rendering is used: + +.. code-block:: python -DataFusion provides extensive customization options for HTML table rendering through the -``datafusion.html_formatter`` module. + # Plain text table output + print(df) -Configuring the HTML Formatter -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Configuring the Formatter +------------------------- -You can customize how DataFrames are rendered by configuring the formatter: +You can customize how DataFrames are rendered by configuring the global formatter: .. code-block:: python - from datafusion.html_formatter import configure_formatter - - # Change the default styling + from datafusion.dataframe_formatter import configure_formatter + configure_formatter( - max_cell_length=25, # Maximum characters in a cell before truncation - max_width=1000, # Maximum width in pixels - max_height=300, # Maximum height in pixels - max_memory_bytes=2097152, # Maximum memory for rendering (2MB) - min_rows=10, # Minimum number of rows to display - max_rows=10, # Maximum rows to display in __repr__ - enable_cell_expansion=True,# Allow expanding truncated cells - custom_css=None, # Additional custom CSS + max_cell_length=25, # Maximum characters in a cell before truncation + max_width=1000, # Maximum width in pixels (HTML only) + max_height=300, # Maximum height in pixels (HTML only) + max_memory_bytes=2097152, # Maximum memory for rendering (2MB) + min_rows=10, # Minimum number of rows to display + max_rows=10, # Maximum rows to display + enable_cell_expansion=True, # Allow expanding truncated cells (HTML only) + custom_css=None, # Additional custom CSS (HTML only) show_truncation_message=True, # Show message when data is truncated - style_provider=None, # Custom styling provider - use_shared_styles=True # Share styles across tables + style_provider=None, # Custom styling provider (HTML only) + use_shared_styles=True, # Share styles across tables (HTML only) ) The formatter settings affect all DataFrames displayed after configuration. Custom Style Providers ------------------------ +---------------------- -For advanced styling needs, you can create a custom style provider: +For HTML styling, you can create a custom style provider that implements the +``StyleProvider`` protocol: .. code-block:: python - from datafusion.html_formatter import StyleProvider, configure_formatter - - class MyStyleProvider(StyleProvider): - def get_table_styles(self): - return { - "table": "border-collapse: collapse; width: 100%;", - "th": "background-color: #007bff; color: white; padding: 8px; text-align: left;", - "td": "border: 1px solid #ddd; padding: 8px;", - "tr:nth-child(even)": "background-color: #f2f2f2;", - } - - def get_value_styles(self, dtype, value): - """Return custom styles for specific values""" - if dtype == "float" and value < 0: - return "color: red;" - return None - + from datafusion.dataframe_formatter import configure_formatter + + class MyStyleProvider: + def get_cell_style(self): + """Return CSS style string for table data cells.""" + return "border: 1px solid #ddd; padding: 8px; text-align: left;" + + def get_header_style(self): + """Return CSS style string for table header cells.""" + return ( + "background-color: #007bff; color: white; " + "padding: 8px; text-align: left;" + ) + # Apply the custom style provider configure_formatter(style_provider=MyStyleProvider()) +Custom Cell Formatters +---------------------- + +You can register custom formatters for specific Python types. A cell formatter is any +callable that takes a value and returns a string: + +.. code-block:: python + + from datafusion.dataframe_formatter import get_formatter + + formatter = get_formatter() + + # Format floats to 2 decimal places + formatter.register_formatter(float, lambda v: f"{v:.2f}") + + # Format dates in a custom way + from datetime import date + formatter.register_formatter(date, lambda v: v.strftime("%B %d, %Y")) + +Custom Cell and Header Builders +------------------------------- + +For full control over the HTML of individual cells or headers, you can set custom +builder functions: + +.. code-block:: python + + from datafusion.dataframe_formatter import get_formatter + + formatter = get_formatter() + + # Custom cell builder receives (value, row, col, table_id) and returns HTML + def my_cell_builder(value, row, col, table_id): + color = "red" if isinstance(value, (int, float)) and value < 0 else "black" + return f"{value}" + + formatter.set_custom_cell_builder(my_cell_builder) + + # Custom header builder receives a schema field and returns HTML + def my_header_builder(field): + return f"{field.name}" + + formatter.set_custom_header_builder(my_header_builder) + Performance Optimization with Shared Styles -------------------------------------------- -The ``use_shared_styles`` parameter (enabled by default) optimizes performance when displaying -multiple DataFrames in notebook environments: +The ``use_shared_styles`` parameter (enabled by default) optimizes performance when +displaying multiple DataFrames in notebook environments: .. code-block:: python - from datafusion.html_formatter import StyleProvider, configure_formatter + from datafusion.dataframe_formatter import configure_formatter + # Default: Use shared styles (recommended for notebooks) configure_formatter(use_shared_styles=True) @@ -111,76 +154,48 @@ multiple DataFrames in notebook environments: configure_formatter(use_shared_styles=False) When ``use_shared_styles=True``: + - CSS styles and JavaScript are included only once per notebook session - This reduces HTML output size and prevents style duplication - Improves rendering performance with many DataFrames - Applies consistent styling across all DataFrames -Creating a Custom Formatter ----------------------------- +Working with the Formatter Directly +------------------------------------ -For complete control over rendering, you can implement a custom formatter: +You can use ``get_formatter()`` and ``set_formatter()`` for direct access to the global +formatter instance: .. code-block:: python - from datafusion.html_formatter import Formatter, get_formatter - - class MyFormatter(Formatter): - def format_html(self, batches, schema, has_more=False, table_uuid=None): - # Create your custom HTML here - html = "
" - # ... formatting logic ... - html += "
" - return html - - # Set as the global formatter - configure_formatter(formatter_class=MyFormatter) - - # Or use the formatter just for specific operations + from datafusion.dataframe_formatter import ( + DataFrameHtmlFormatter, + get_formatter, + set_formatter, + ) + + # Get and modify the current formatter formatter = get_formatter() - custom_html = formatter.format_html(batches, schema) + print(formatter.max_rows) + print(formatter.max_cell_length) -Managing Formatters -------------------- + # Create and set a fully custom formatter + custom_formatter = DataFrameHtmlFormatter( + max_cell_length=50, + max_rows=20, + enable_cell_expansion=False, + ) + set_formatter(custom_formatter) Reset to default formatting: .. code-block:: python - from datafusion.html_formatter import reset_formatter - + from datafusion.dataframe_formatter import reset_formatter + # Reset to default settings reset_formatter() -Get the current formatter settings: - -.. code-block:: python - - from datafusion.html_formatter import get_formatter - - formatter = get_formatter() - print(formatter.max_rows) - print(formatter.theme) - -Contextual Formatting ----------------------- - -You can also use a context manager to temporarily change formatting settings: - -.. code-block:: python - - from datafusion.html_formatter import formatting_context - - # Default formatting - df.show() - - # Temporarily use different formatting - with formatting_context(max_rows=100, theme="dark"): - df.show() # Will use the temporary settings - - # Back to default formatting - df.show() - Memory and Display Controls --------------------------- @@ -188,10 +203,12 @@ You can control how much data is displayed and how much memory is used for rende .. code-block:: python + from datafusion.dataframe_formatter import configure_formatter + configure_formatter( max_memory_bytes=4 * 1024 * 1024, # 4MB maximum memory for display min_rows=20, # Always show at least 20 rows - max_rows=50 # Show up to 50 rows in output + max_rows=50, # Show up to 50 rows in output ) These parameters help balance comprehensive data display against performance considerations. @@ -216,7 +233,7 @@ Additional Resources * :doc:`../io/index` - I/O Guide for reading data from various sources * :doc:`../data-sources` - Comprehensive data sources guide * :ref:`io_csv` - CSV file reading -* :ref:`io_parquet` - Parquet file reading +* :ref:`io_parquet` - Parquet file reading * :ref:`io_json` - JSON file reading * :ref:`io_avro` - Avro file reading * :ref:`io_custom_table_provider` - Custom table providers diff --git a/examples/tpch/q02_minimum_cost_supplier.py b/examples/tpch/q02_minimum_cost_supplier.py index 7390d0892..47961d2ef 100644 --- a/examples/tpch/q02_minimum_cost_supplier.py +++ b/examples/tpch/q02_minimum_cost_supplier.py @@ -32,6 +32,7 @@ import datafusion from datafusion import SessionContext, col, lit from datafusion import functions as F +from datafusion.expr import Window from util import get_data_path # This is the part we're looking for. Values selected here differ from the spec in order to run @@ -106,11 +107,8 @@ window_frame = datafusion.WindowFrame("rows", None, None) df = df.with_column( "min_cost", - F.window( - "min", - [col("ps_supplycost")], - partition_by=[col("ps_partkey")], - window_frame=window_frame, + F.min(col("ps_supplycost")).over( + Window(partition_by=[col("ps_partkey")], window_frame=window_frame) ), ) diff --git a/examples/tpch/q11_important_stock_identification.py b/examples/tpch/q11_important_stock_identification.py index 22829ab7c..de309fa64 100644 --- a/examples/tpch/q11_important_stock_identification.py +++ b/examples/tpch/q11_important_stock_identification.py @@ -29,6 +29,7 @@ from datafusion import SessionContext, WindowFrame, col, lit from datafusion import functions as F +from datafusion.expr import Window from util import get_data_path NATION = "GERMANY" @@ -71,7 +72,7 @@ window_frame = WindowFrame("rows", None, None) df = df.with_column( - "total_value", F.window("sum", [col("value")], window_frame=window_frame) + "total_value", F.sum(col("value")).over(Window(window_frame=window_frame)) ) # Limit to the parts for which there is a significant value based on the fraction of the total diff --git a/examples/tpch/q15_top_supplier.py b/examples/tpch/q15_top_supplier.py index c321048f2..5128937a7 100644 --- a/examples/tpch/q15_top_supplier.py +++ b/examples/tpch/q15_top_supplier.py @@ -31,6 +31,7 @@ import pyarrow as pa from datafusion import SessionContext, WindowFrame, col, lit from datafusion import functions as F +from datafusion.expr import Window from util import get_data_path DATE = "1996-01-01" @@ -70,7 +71,8 @@ # Use a window function to find the maximum revenue across the entire dataframe window_frame = WindowFrame("rows", None, None) df = df.with_column( - "max_revenue", F.window("max", [col("total_revenue")], window_frame=window_frame) + "max_revenue", + F.max(col("total_revenue")).over(Window(window_frame=window_frame)), ) # Find all suppliers whose total revenue is the same as the maximum diff --git a/examples/tpch/q17_small_quantity_order.py b/examples/tpch/q17_small_quantity_order.py index 6d76fe506..5ccb38422 100644 --- a/examples/tpch/q17_small_quantity_order.py +++ b/examples/tpch/q17_small_quantity_order.py @@ -30,6 +30,7 @@ from datafusion import SessionContext, WindowFrame, col, lit from datafusion import functions as F +from datafusion.expr import Window from util import get_data_path BRAND = "Brand#23" @@ -58,11 +59,8 @@ window_frame = WindowFrame("rows", None, None) df = df.with_column( "avg_quantity", - F.window( - "avg", - [col("l_quantity")], - window_frame=window_frame, - partition_by=[col("l_partkey")], + F.avg(col("l_quantity")).over( + Window(partition_by=[col("l_partkey")], window_frame=window_frame) ), ) diff --git a/examples/tpch/q22_global_sales_opportunity.py b/examples/tpch/q22_global_sales_opportunity.py index c4d115b74..a2d41b215 100644 --- a/examples/tpch/q22_global_sales_opportunity.py +++ b/examples/tpch/q22_global_sales_opportunity.py @@ -28,6 +28,7 @@ from datafusion import SessionContext, WindowFrame, col, lit from datafusion import functions as F +from datafusion.expr import Window from util import get_data_path NATION_CODES = [13, 31, 23, 29, 30, 18, 17] @@ -55,7 +56,8 @@ # current row. We want our frame to cover the entire data frame. window_frame = WindowFrame("rows", None, None) df = df.with_column( - "avg_balance", F.window("avg", [col("c_acctbal")], window_frame=window_frame) + "avg_balance", + F.avg(col("c_acctbal")).over(Window(window_frame=window_frame)), ) df.show() diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index a736c3966..ee02c921d 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -35,7 +35,7 @@ # The following imports are okay to remain as opaque to the user. from ._internal import Config -from .catalog import Catalog, Database, Table +from .catalog import Catalog, Table from .col import col, column from .common import DFSchema from .context import ( @@ -81,7 +81,6 @@ "DFSchema", "DataFrame", "DataFrameWriteOptions", - "Database", "ExecutionPlan", "ExplainFormat", "Expr", diff --git a/python/datafusion/catalog.py b/python/datafusion/catalog.py index 03c0ddc68..20da5e671 100644 --- a/python/datafusion/catalog.py +++ b/python/datafusion/catalog.py @@ -129,11 +129,6 @@ def schema(self, name: str = "public") -> Schema: else schema ) - @deprecated("Use `schema` instead.") - def database(self, name: str = "public") -> Schema: - """Returns the database with the given ``name`` from this catalog.""" - return self.schema(name) - def register_schema( self, name: str, @@ -195,11 +190,6 @@ def table_exist(self, name: str) -> bool: return self._raw_schema.table_exist(name) -@deprecated("Use `Schema` instead.") -class Database(Schema): - """See `Schema`.""" - - class Table: """A DataFusion table. diff --git a/python/datafusion/context.py b/python/datafusion/context.py index e3949de83..c3f94cc16 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -426,11 +426,6 @@ def with_temp_file_path(self, path: str | pathlib.Path) -> RuntimeEnvBuilder: return self -@deprecated("Use `RuntimeEnvBuilder` instead.") -class RuntimeConfig(RuntimeEnvBuilder): - """See `RuntimeEnvBuilder`.""" - - class SQLOptions: """Options to be used when performing SQL queries.""" @@ -785,14 +780,6 @@ def from_arrow( """ return DataFrame(self.ctx.from_arrow(data, name)) - @deprecated("Use ``from_arrow`` instead.") - def from_arrow_table(self, data: pa.Table, name: str | None = None) -> DataFrame: - """Create a :py:class:`~datafusion.dataframe.DataFrame` from an Arrow table. - - This is an alias for :py:func:`from_arrow`. - """ - return self.from_arrow(data, name) - def from_pandas(self, data: pd.DataFrame, name: str | None = None) -> DataFrame: """Create a :py:class:`~datafusion.dataframe.DataFrame` from a Pandas DataFrame. @@ -1260,14 +1247,6 @@ def catalog(self, name: str = "datafusion") -> Catalog: """Retrieve a catalog by name.""" return Catalog(self.ctx.catalog(name)) - @deprecated( - "Use the catalog provider interface ``SessionContext.Catalog`` to " - "examine available catalogs, schemas and tables" - ) - def tables(self) -> set[str]: - """Deprecated.""" - return self.ctx.tables() - def table(self, name: str) -> DataFrame: """Retrieve a previously registered table by name.""" return DataFrame(self.ctx.table(name)) diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index 9dc5f0e7d..c00c85fdb 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -489,17 +489,6 @@ def find_qualified_columns(self, *names: str) -> list[Expr]: raw_exprs = self.df.find_qualified_columns(list(names)) return [Expr(e) for e in raw_exprs] - @deprecated( - "select_columns() is deprecated. Use :py:meth:`~DataFrame.select` instead" - ) - def select_columns(self, *args: str) -> DataFrame: - """Filter the DataFrame by columns. - - Returns: - DataFrame only containing the specified columns. - """ - return self.select(*args) - def select_exprs(self, *args: str) -> DataFrame: """Project arbitrary list of expression strings into a new DataFrame. @@ -1603,15 +1592,6 @@ def count(self) -> int: """ return self.df.count() - @deprecated("Use :py:func:`unnest_columns` instead.") - def unnest_column( - self, - column: str, - preserve_nulls: bool = True, - ) -> DataFrame: - """See :py:func:`unnest_columns`.""" - return DataFrame(self.df.unnest_column(column, preserve_nulls=preserve_nulls)) - def unnest_columns( self, *columns: str, diff --git a/python/datafusion/dataframe_formatter.py b/python/datafusion/dataframe_formatter.py index b8af45a1b..fd2da99f0 100644 --- a/python/datafusion/dataframe_formatter.py +++ b/python/datafusion/dataframe_formatter.py @@ -748,7 +748,7 @@ def get_formatter() -> DataFrameHtmlFormatter: The global HTML formatter instance Example: - >>> from datafusion.html_formatter import get_formatter + >>> from datafusion.dataframe_formatter import get_formatter >>> formatter = get_formatter() >>> formatter.max_cell_length = 50 # Increase cell length """ @@ -762,7 +762,7 @@ def set_formatter(formatter: DataFrameHtmlFormatter) -> None: formatter: The formatter instance to use globally Example: - >>> from datafusion.html_formatter import get_formatter, set_formatter + >>> from datafusion.dataframe_formatter import get_formatter, set_formatter >>> custom_formatter = DataFrameHtmlFormatter(max_cell_length=100) >>> set_formatter(custom_formatter) """ @@ -783,7 +783,7 @@ def configure_formatter(**kwargs: Any) -> None: ValueError: If any invalid parameters are provided Example: - >>> from datafusion.html_formatter import configure_formatter + >>> from datafusion.dataframe_formatter import configure_formatter >>> configure_formatter( ... max_cell_length=50, ... max_height=500, @@ -827,7 +827,7 @@ def reset_formatter() -> None: and sets it as the global formatter for all DataFrames. Example: - >>> from datafusion.html_formatter import reset_formatter + >>> from datafusion.dataframe_formatter import reset_formatter >>> reset_formatter() # Reset formatter to default settings """ formatter = DataFrameHtmlFormatter() diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 35388468c..7cd74ecd5 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -27,11 +27,6 @@ from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING, Any, ClassVar -try: - from warnings import deprecated # Python 3.13+ -except ImportError: - from typing_extensions import deprecated # Python 3.12 - import pyarrow as pa from ._internal import expr as expr_internal @@ -356,16 +351,6 @@ def to_variant(self) -> Any: """Convert this expression into a python object if possible.""" return self.expr.to_variant() - @deprecated( - "display_name() is deprecated. Use :py:meth:`~Expr.schema_name` instead" - ) - def display_name(self) -> str: - """Returns the name of this expression as it should appear in a schema. - - This name will not include any CAST expressions. - """ - return self.schema_name() - def schema_name(self) -> str: """Returns the name of this expression as it should appear in a schema. diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 9dfabb62d..841cd9c0b 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any import pyarrow as pa @@ -29,19 +29,11 @@ Expr, SortExpr, SortKey, - WindowFrame, expr_list_to_raw_expr_list, sort_list_to_raw_sort_list, sort_or_default, ) -try: - from warnings import deprecated # Python 3.13+ -except ImportError: - from typing_extensions import deprecated # Python 3.12 - -if TYPE_CHECKING: - from datafusion.context import SessionContext __all__ = [ "abs", "acos", @@ -339,8 +331,6 @@ "var_sample", "version", "when", - # Window Functions - "window", ] @@ -664,49 +654,6 @@ def when(when: Expr, then: Expr) -> CaseBuilder: return CaseBuilder(f.when(when.expr, then.expr)) -@deprecated("Prefer to call Expr.over() instead") -def window( - name: str, - args: list[Expr], - partition_by: list[Expr] | Expr | None = None, - order_by: list[SortKey] | SortKey | None = None, - window_frame: WindowFrame | None = None, - filter: Expr | None = None, - distinct: bool = False, - ctx: SessionContext | None = None, -) -> Expr: - """Creates a new Window function expression. - - This interface will soon be deprecated. Instead of using this interface, - users should call the window functions directly. For example, to perform a - lag use:: - - df.select(functions.lag(col("a")).partition_by(col("b")).build()) - - The ``order_by`` parameter accepts column names or expressions, e.g.:: - - window("lag", [col("a")], order_by="ts") - """ - args = [a.expr for a in args] - partition_by_raw = expr_list_to_raw_expr_list(partition_by) - order_by_raw = sort_list_to_raw_sort_list(order_by) - window_frame = window_frame.window_frame if window_frame is not None else None - ctx = ctx.ctx if ctx is not None else None - filter_raw = filter.expr if filter is not None else None - return Expr( - f.window( - name, - args, - partition_by=partition_by_raw, - order_by=order_by_raw, - window_frame=window_frame, - ctx=ctx, - filter=filter_raw, - distinct=distinct, - ) - ) - - # scalar functions def abs(arg: Expr) -> Expr: """Return the absolute value of a given number. diff --git a/python/datafusion/html_formatter.py b/python/datafusion/html_formatter.py deleted file mode 100644 index 65eb1f042..000000000 --- a/python/datafusion/html_formatter.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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. - -"""Deprecated module for dataframe formatting.""" - -import warnings - -from datafusion.dataframe_formatter import * # noqa: F403 - -warnings.warn( - "The module 'html_formatter' is deprecated and will be removed in the next release." - "Please use 'dataframe_formatter' instead.", - DeprecationWarning, - stacklevel=3, -) diff --git a/python/datafusion/substrait.py b/python/datafusion/substrait.py index 3115238fa..6353ef8cc 100644 --- a/python/datafusion/substrait.py +++ b/python/datafusion/substrait.py @@ -25,11 +25,6 @@ from typing import TYPE_CHECKING -try: - from warnings import deprecated # Python 3.13+ -except ImportError: - from typing_extensions import deprecated # Python 3.12 - from datafusion.plan import LogicalPlan from ._internal import substrait as substrait_internal @@ -88,11 +83,6 @@ def from_json(json: str) -> Plan: return Plan(substrait_internal.Plan.from_json(json)) -@deprecated("Use `Plan` instead.") -class plan(Plan): # noqa: N801 - """See `Plan`.""" - - class Serde: """Provides the ``Substrait`` serialization and deserialization.""" @@ -158,11 +148,6 @@ def deserialize_bytes(proto_bytes: bytes) -> Plan: return Plan(substrait_internal.Serde.deserialize_bytes(proto_bytes)) -@deprecated("Use `Serde` instead.") -class serde(Serde): # noqa: N801 - """See `Serde` instead.""" - - class Producer: """Generates substrait plans from a logical plan.""" @@ -184,11 +169,6 @@ def to_substrait_plan(logical_plan: LogicalPlan, ctx: SessionContext) -> Plan: ) -@deprecated("Use `Producer` instead.") -class producer(Producer): # noqa: N801 - """Use `Producer` instead.""" - - class Consumer: """Generates a logical plan from a substrait plan.""" @@ -206,8 +186,3 @@ def from_substrait_plan(ctx: SessionContext, plan: Plan) -> LogicalPlan: return LogicalPlan( substrait_internal.Consumer.from_substrait_plan(ctx.ctx, plan.plan_internal) ) - - -@deprecated("Use `Consumer` instead.") -class consumer(Consumer): # noqa: N801 - """Use `Consumer` instead.""" diff --git a/python/datafusion/udf.py b/python/datafusion/udf.py deleted file mode 100644 index c7265fa09..000000000 --- a/python/datafusion/udf.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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. - -"""Deprecated module for user defined functions.""" - -import warnings - -from datafusion.user_defined import * # noqa: F403 - -warnings.warn( - "The module 'udf' is deprecated and will be removed in the next release. " - "Please use 'user_defined' instead.", - DeprecationWarning, - stacklevel=2, -) diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index 9a287c1f7..1cf824a15 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -319,27 +319,6 @@ def test_expr_getitem() -> None: assert array_values == [2, 5, None, None] -def test_display_name_deprecation(): - import warnings - - expr = col("foo") - with warnings.catch_warnings(record=True) as w: - # Cause all warnings to always be triggered - warnings.simplefilter("always") - - # should trigger warning - name = expr.display_name() - - # Verify some things - assert len(w) == 1 - assert issubclass(w[-1].category, DeprecationWarning) - assert "deprecated" in str(w[-1].message) - - # returns appropriate result - assert name == expr.schema_name() - assert name == "foo" - - @pytest.fixture def df(): ctx = SessionContext() From 1be838bb47f04bcf4d1a0f65e3e6958aa9366f3f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sun, 12 Apr 2026 21:24:39 -0400 Subject: [PATCH 18/83] Release 53.0.0 (#1491) * Update version number and changelog * minor: set version number on dependency to publish to crates.io * taplo fmt --- Cargo.lock | 6 +-- Cargo.toml | 6 +-- dev/changelog/53.0.0.md | 107 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 6 deletions(-) create mode 100644 dev/changelog/53.0.0.md diff --git a/Cargo.lock b/Cargo.lock index ee89c8bda..1cbb0acb8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1312,7 +1312,7 @@ dependencies = [ [[package]] name = "datafusion-ffi-example" -version = "52.0.0" +version = "53.0.0" dependencies = [ "arrow", "arrow-array", @@ -1662,7 +1662,7 @@ dependencies = [ [[package]] name = "datafusion-python" -version = "52.0.0" +version = "53.0.0" dependencies = [ "arrow", "arrow-select", @@ -1692,7 +1692,7 @@ dependencies = [ [[package]] name = "datafusion-python-util" -version = "52.0.0" +version = "53.0.0" dependencies = [ "arrow", "datafusion", diff --git a/Cargo.toml b/Cargo.toml index 3a34e204c..14408d2bc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ # under the License. [workspace.package] -version = "52.0.0" +version = "53.0.0" homepage = "https://datafusion.apache.org/python" repository = "https://github.com/apache/datafusion-python" authors = ["Apache DataFusion "] @@ -59,9 +59,9 @@ object_store = { version = "0.13.1" } url = "2" log = "0.4.29" parking_lot = "0.12" -prost-types = "0.14.3" # keep in line with `datafusion-substrait` +prost-types = "0.14.3" # keep in line with `datafusion-substrait` pyo3-build-config = "0.28" -datafusion-python-util = { path = "crates/util" } +datafusion-python-util = { path = "crates/util", version = "53.0.0" } [profile.release] lto = "thin" diff --git a/dev/changelog/53.0.0.md b/dev/changelog/53.0.0.md new file mode 100644 index 000000000..3e27a852d --- /dev/null +++ b/dev/changelog/53.0.0.md @@ -0,0 +1,107 @@ + + +# Apache DataFusion Python 53.0.0 Changelog + +This release consists of 52 commits from 9 contributors. See credits at the end of this changelog for more information. + +**Breaking changes:** + +- minor: remove deprecated interfaces [#1481](https://github.com/apache/datafusion-python/pull/1481) (timsaucer) + +**Implemented enhancements:** + +- feat: feat: add to_time, to_local_time, to_date functions [#1387](https://github.com/apache/datafusion-python/pull/1387) (mesejo) +- feat: Add FFI_TableProviderFactory support [#1396](https://github.com/apache/datafusion-python/pull/1396) (davisp) + +**Fixed bugs:** + +- fix: satisfy rustfmt check in lib.rs re-exports [#1406](https://github.com/apache/datafusion-python/pull/1406) (kevinjqliu) + +**Documentation updates:** + +- docs: clarify DataFusion 52 FFI session-parameter requirement for provider hooks [#1439](https://github.com/apache/datafusion-python/pull/1439) (kevinjqliu) + +**Other:** + +- Merge release 52.0.0 into main [#1389](https://github.com/apache/datafusion-python/pull/1389) (timsaucer) +- Add workflow to verify release candidate on multiple systems [#1388](https://github.com/apache/datafusion-python/pull/1388) (timsaucer) +- Allow running "verify release candidate" github workflow on Windows [#1392](https://github.com/apache/datafusion-python/pull/1392) (kevinjqliu) +- ci: update pre-commit hooks, fix linting, and refresh dependencies [#1385](https://github.com/apache/datafusion-python/pull/1385) (dariocurr) +- Add CI check for crates.io patches [#1407](https://github.com/apache/datafusion-python/pull/1407) (timsaucer) +- Enable doc tests in local and CI testing [#1409](https://github.com/apache/datafusion-python/pull/1409) (ntjohnson1) +- Upgrade to DataFusion 53 [#1402](https://github.com/apache/datafusion-python/pull/1402) (nuno-faria) +- Catch warnings in FFI unit tests [#1410](https://github.com/apache/datafusion-python/pull/1410) (timsaucer) +- Add docstring examples for Scalar trigonometric functions [#1411](https://github.com/apache/datafusion-python/pull/1411) (ntjohnson1) +- Create workspace with core and util crates [#1414](https://github.com/apache/datafusion-python/pull/1414) (timsaucer) +- Add docstring examples for Scalar regex, crypto, struct and other [#1422](https://github.com/apache/datafusion-python/pull/1422) (ntjohnson1) +- Add docstring examples for Scalar math functions [#1421](https://github.com/apache/datafusion-python/pull/1421) (ntjohnson1) +- Add docstring examples for Common utility functions [#1419](https://github.com/apache/datafusion-python/pull/1419) (ntjohnson1) +- Add docstring examples for Aggregate basic and bitwise/boolean functions [#1416](https://github.com/apache/datafusion-python/pull/1416) (ntjohnson1) +- Fix CI errors on main [#1432](https://github.com/apache/datafusion-python/pull/1432) (timsaucer) +- Add docstring examples for Scalar temporal functions [#1424](https://github.com/apache/datafusion-python/pull/1424) (ntjohnson1) +- Add docstring examples for Aggregate statistical and regression functions [#1417](https://github.com/apache/datafusion-python/pull/1417) (ntjohnson1) +- Add docstring examples for Scalar array/list functions [#1420](https://github.com/apache/datafusion-python/pull/1420) (ntjohnson1) +- Add docstring examples for Scalar string functions [#1423](https://github.com/apache/datafusion-python/pull/1423) (ntjohnson1) +- Add docstring examples for Aggregate window functions [#1418](https://github.com/apache/datafusion-python/pull/1418) (ntjohnson1) +- ci: pin third-party actions to Apache-approved SHAs [#1438](https://github.com/apache/datafusion-python/pull/1438) (kevinjqliu) +- minor: bump datafusion to release version [#1441](https://github.com/apache/datafusion-python/pull/1441) (timsaucer) +- ci: add swap during build, use tpchgen-cli [#1443](https://github.com/apache/datafusion-python/pull/1443) (timsaucer) +- Update remaining existing examples to make testable/standalone executable [#1437](https://github.com/apache/datafusion-python/pull/1437) (ntjohnson1) +- Do not run validate_pycapsule if pointer_checked is used [#1426](https://github.com/apache/datafusion-python/pull/1426) (Tpt) +- Implement configuration extension support [#1391](https://github.com/apache/datafusion-python/pull/1391) (timsaucer) +- Add a working, more complete example of using a catalog (docs) [#1427](https://github.com/apache/datafusion-python/pull/1427) (toppyy) +- chore: update dependencies [#1447](https://github.com/apache/datafusion-python/pull/1447) (timsaucer) +- Complete doc string examples for functions.py [#1435](https://github.com/apache/datafusion-python/pull/1435) (ntjohnson1) +- chore: enforce uv lockfile consistency in CI and pre-commit [#1398](https://github.com/apache/datafusion-python/pull/1398) (mesejo) +- CI: Add CodeQL workflow for GitHub Actions security scanning [#1408](https://github.com/apache/datafusion-python/pull/1408) (kevinjqliu) +- ci: update codespell paths [#1469](https://github.com/apache/datafusion-python/pull/1469) (timsaucer) +- Add missing datetime functions [#1467](https://github.com/apache/datafusion-python/pull/1467) (timsaucer) +- Add AI skill to check current repository against upstream APIs [#1460](https://github.com/apache/datafusion-python/pull/1460) (timsaucer) +- Add missing string function `contains` [#1465](https://github.com/apache/datafusion-python/pull/1465) (timsaucer) +- Add missing conditional functions [#1464](https://github.com/apache/datafusion-python/pull/1464) (timsaucer) +- Reduce peak memory usage during release builds to fix OOM on manylinux runners [#1445](https://github.com/apache/datafusion-python/pull/1445) (kevinjqliu) +- Add missing map functions [#1461](https://github.com/apache/datafusion-python/pull/1461) (timsaucer) +- minor: Fix pytest instructions in the README [#1477](https://github.com/apache/datafusion-python/pull/1477) (nuno-faria) +- Add missing array functions [#1468](https://github.com/apache/datafusion-python/pull/1468) (timsaucer) +- Add missing scalar functions [#1470](https://github.com/apache/datafusion-python/pull/1470) (timsaucer) +- Add missing aggregate functions [#1471](https://github.com/apache/datafusion-python/pull/1471) (timsaucer) +- Add missing Dataframe functions [#1472](https://github.com/apache/datafusion-python/pull/1472) (timsaucer) +- Add missing deregister methods to SessionContext [#1473](https://github.com/apache/datafusion-python/pull/1473) (timsaucer) +- Add missing registration methods [#1474](https://github.com/apache/datafusion-python/pull/1474) (timsaucer) +- Add missing SessionContext utility methods [#1475](https://github.com/apache/datafusion-python/pull/1475) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 25 Tim Saucer + 13 Nick + 6 Kevin Liu + 2 Daniel Mesejo + 2 Nuno Faria + 1 Paul J. Davis + 1 Thomas Tanon + 1 Topias Pyykkönen + 1 dario curreri +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + From 00b24572c98a257f06ff026a90c07634a86204d4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 13 Apr 2026 06:33:29 -0400 Subject: [PATCH 19/83] ci: disable symbol export on Windows verification (#1486) * Set rust flags on windows release verification * Forward flag to linker * Switch to msvc rust toolchain * Revert "Switch to msvc rust toolchain" This reverts commit 9879fc7dbe066098445b9600087e665435b58f8a. --- .github/workflows/verify-release-candidate.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/verify-release-candidate.yml b/.github/workflows/verify-release-candidate.yml index a10a4faa9..7a5deff5b 100644 --- a/.github/workflows/verify-release-candidate.yml +++ b/.github/workflows/verify-release-candidate.yml @@ -73,6 +73,11 @@ jobs: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Set RUSTFLAGS for Windows GNU linker + if: matrix.os == 'windows' + shell: bash + run: echo "RUSTFLAGS=-C link-arg=-Wl,--exclude-libs=ALL" >> "$GITHUB_ENV" + - name: Run release candidate verification shell: bash run: ./dev/release/verify-release-candidate.sh "${{ inputs.version }}" "${{ inputs.rc_number }}" From 8a7efead43cff8dc7515e27e53da7545100e25a7 Mon Sep 17 00:00:00 2001 From: Shreyesh Date: Mon, 13 Apr 2026 03:34:35 -0700 Subject: [PATCH 20/83] Add Python bindings for accessing ExecutionMetrics (#1381) * feat: add Python bindings for accessing ExecutionMetrics * test: imporve tests * first round of reviews * plan caching * address some concerns * merge and address comments * fix Ci issues * attempt to fix lint * fix build * fix docstring * address some more comments --------- Co-authored-by: ShreyeshArangath --- Cargo.lock | 1 + Cargo.toml | 1 + crates/core/Cargo.toml | 1 + crates/core/src/dataframe.rs | 49 +++- crates/core/src/lib.rs | 3 + crates/core/src/metrics.rs | 169 ++++++++++++++ crates/core/src/physical_plan.rs | 5 + .../dataframe/execution-metrics.rst | 215 ++++++++++++++++++ docs/source/user-guide/dataframe/index.rst | 9 + python/datafusion/__init__.py | 4 +- python/datafusion/plan.py | 177 ++++++++++++++ python/tests/test_plans.py | 192 +++++++++++++++- 12 files changed, 817 insertions(+), 9 deletions(-) create mode 100644 crates/core/src/metrics.rs create mode 100644 docs/source/user-guide/dataframe/execution-metrics.rst diff --git a/Cargo.lock b/Cargo.lock index 1cbb0acb8..4efca3eb6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1667,6 +1667,7 @@ dependencies = [ "arrow", "arrow-select", "async-trait", + "chrono", "cstr", "datafusion", "datafusion-ffi", diff --git a/Cargo.toml b/Cargo.toml index 14408d2bc..d0e87a9a4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ tokio = { version = "1.50" } pyo3 = { version = "0.28" } pyo3-async-runtimes = { version = "0.28" } pyo3-log = "0.13.3" +chrono = { version = "0.4", default-features = false } arrow = { version = "58" } arrow-array = { version = "58" } arrow-schema = { version = "58" } diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 3e2b01c8e..d714dc978 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -47,6 +47,7 @@ pyo3 = { workspace = true, features = [ ] } pyo3-async-runtimes = { workspace = true, features = ["tokio-runtime"] } pyo3-log = { workspace = true } +chrono = { workspace = true } arrow = { workspace = true, features = ["pyarrow"] } arrow-select = { workspace = true } datafusion = { workspace = true, features = ["avro", "unicode_expressions"] } diff --git a/crates/core/src/dataframe.rs b/crates/core/src/dataframe.rs index c067eac30..2d815ec76 100644 --- a/crates/core/src/dataframe.rs +++ b/crates/core/src/dataframe.rs @@ -37,9 +37,15 @@ use datafusion::config::{CsvOptions, ParquetColumnOptions, ParquetOptions, Table use datafusion::dataframe::{DataFrame, DataFrameWriteOptions}; use datafusion::error::DataFusionError; use datafusion::execution::SendableRecordBatchStream; +use datafusion::execution::context::TaskContext; use datafusion::logical_expr::SortExpr; use datafusion::logical_expr::dml::InsertOp; use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; +use datafusion::physical_plan::{ + ExecutionPlan as DFExecutionPlan, collect as df_collect, + collect_partitioned as df_collect_partitioned, execute_stream as df_execute_stream, + execute_stream_partitioned as df_execute_stream_partitioned, +}; use datafusion::prelude::*; use datafusion_python_util::{is_ipython_env, spawn_future, wait_for_future}; use futures::{StreamExt, TryStreamExt}; @@ -308,6 +314,9 @@ pub struct PyDataFrame { // In IPython environment cache batches between __repr__ and _repr_html_ calls. batches: SharedCachedBatches, + + // Cache the last physical plan so that metrics are available after execution. + last_plan: Arc>>>, } impl PyDataFrame { @@ -316,6 +325,7 @@ impl PyDataFrame { Self { df: Arc::new(df), batches: Arc::new(Mutex::new(None)), + last_plan: Arc::new(Mutex::new(None)), } } @@ -387,6 +397,20 @@ impl PyDataFrame { Ok(html_str) } + /// Create the physical plan, cache it in `last_plan`, and return the plan together + /// with a task context. Centralises the repeated three-line pattern that appears in + /// `collect`, `collect_partitioned`, `execute_stream`, and `execute_stream_partitioned`. + fn create_and_cache_plan( + &self, + py: Python, + ) -> PyDataFusionResult<(Arc, Arc)> { + let df = self.df.as_ref().clone(); + let new_plan = wait_for_future(py, df.create_physical_plan())??; + *self.last_plan.lock() = Some(Arc::clone(&new_plan)); + let task_ctx = Arc::new(self.df.as_ref().task_ctx()); + Ok((new_plan, task_ctx)) + } + async fn collect_column_inner(&self, column: &str) -> Result { let batches = self .df @@ -646,8 +670,9 @@ impl PyDataFrame { /// Unless some order is specified in the plan, there is no /// guarantee of the order of the result. fn collect<'py>(&self, py: Python<'py>) -> PyResult>> { - let batches = wait_for_future(py, self.df.as_ref().clone().collect())? - .map_err(PyDataFusionError::from)?; + let (plan, task_ctx) = self.create_and_cache_plan(py)?; + let batches = + wait_for_future(py, df_collect(plan, task_ctx))?.map_err(PyDataFusionError::from)?; // cannot use PyResult> return type due to // https://github.com/PyO3/pyo3/issues/1813 batches.into_iter().map(|rb| rb.to_pyarrow(py)).collect() @@ -662,7 +687,8 @@ impl PyDataFrame { /// Executes this DataFrame and collects all results into a vector of vector of RecordBatch /// maintaining the input partitioning. fn collect_partitioned<'py>(&self, py: Python<'py>) -> PyResult>>> { - let batches = wait_for_future(py, self.df.as_ref().clone().collect_partitioned())? + let (plan, task_ctx) = self.create_and_cache_plan(py)?; + let batches = wait_for_future(py, df_collect_partitioned(plan, task_ctx))? .map_err(PyDataFusionError::from)?; batches @@ -840,7 +866,13 @@ impl PyDataFrame { } /// Get the execution plan for this `DataFrame` + /// + /// If the DataFrame has already been executed (e.g. via `collect()`), + /// returns the cached plan which includes populated metrics. fn execution_plan(&self, py: Python) -> PyDataFusionResult { + if let Some(plan) = self.last_plan.lock().as_ref() { + return Ok(PyExecutionPlan::new(Arc::clone(plan))); + } let plan = wait_for_future(py, self.df.as_ref().clone().create_physical_plan())??; Ok(plan.into()) } @@ -1198,14 +1230,17 @@ impl PyDataFrame { } fn execute_stream(&self, py: Python) -> PyDataFusionResult { - let df = self.df.as_ref().clone(); - let stream = spawn_future(py, async move { df.execute_stream().await })?; + let (plan, task_ctx) = self.create_and_cache_plan(py)?; + let stream = spawn_future(py, async move { df_execute_stream(plan, task_ctx) })?; Ok(PyRecordBatchStream::new(stream)) } fn execute_stream_partitioned(&self, py: Python) -> PyResult> { - let df = self.df.as_ref().clone(); - let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?; + let (plan, task_ctx) = self.create_and_cache_plan(py)?; + let streams = spawn_future( + py, + async move { df_execute_stream_partitioned(plan, task_ctx) }, + )?; Ok(streams.into_iter().map(PyRecordBatchStream::new).collect()) } diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index fc2d006d3..77d69911a 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -43,6 +43,7 @@ pub mod errors; pub mod expr; #[allow(clippy::borrow_deref_ref)] mod functions; +pub mod metrics; mod options; pub mod physical_plan; mod pyarrow_filter_expression; @@ -92,6 +93,8 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/core/src/metrics.rs b/crates/core/src/metrics.rs new file mode 100644 index 000000000..ee0937e25 --- /dev/null +++ b/crates/core/src/metrics.rs @@ -0,0 +1,169 @@ +// 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. + +use std::collections::HashMap; +use std::sync::Arc; + +use chrono::{Datelike, Timelike}; +use datafusion::physical_plan::metrics::{Metric, MetricValue, MetricsSet, Timestamp}; +use pyo3::prelude::*; + +#[pyclass(from_py_object, frozen, name = "MetricsSet", module = "datafusion")] +#[derive(Debug, Clone)] +pub struct PyMetricsSet { + metrics: MetricsSet, +} + +impl PyMetricsSet { + pub fn new(metrics: MetricsSet) -> Self { + Self { metrics } + } +} + +#[pymethods] +impl PyMetricsSet { + fn metrics(&self) -> Vec { + self.metrics + .iter() + .map(|m| PyMetric::new(Arc::clone(m))) + .collect() + } + + fn output_rows(&self) -> Option { + self.metrics.output_rows() + } + + fn elapsed_compute(&self) -> Option { + self.metrics.elapsed_compute() + } + + fn spill_count(&self) -> Option { + self.metrics.spill_count() + } + + fn spilled_bytes(&self) -> Option { + self.metrics.spilled_bytes() + } + + fn spilled_rows(&self) -> Option { + self.metrics.spilled_rows() + } + + fn sum_by_name(&self, name: &str) -> Option { + self.metrics.sum_by_name(name).map(|v| v.as_usize()) + } + + fn __repr__(&self) -> String { + format!("{}", self.metrics) + } +} + +#[pyclass(from_py_object, frozen, name = "Metric", module = "datafusion")] +#[derive(Debug, Clone)] +pub struct PyMetric { + metric: Arc, +} + +impl PyMetric { + pub fn new(metric: Arc) -> Self { + Self { metric } + } + + fn timestamp_to_pyobject<'py>( + py: Python<'py>, + ts: &Timestamp, + ) -> PyResult>> { + match ts.value() { + Some(dt) => { + let datetime_mod = py.import("datetime")?; + let datetime_cls = datetime_mod.getattr("datetime")?; + let tz_utc = datetime_mod.getattr("timezone")?.getattr("utc")?; + let result = datetime_cls.call1(( + dt.year(), + dt.month(), + dt.day(), + dt.hour(), + dt.minute(), + dt.second(), + dt.timestamp_subsec_micros(), + tz_utc, + ))?; + Ok(Some(result)) + } + None => Ok(None), + } + } +} + +#[pymethods] +impl PyMetric { + #[getter] + fn name(&self) -> String { + self.metric.value().name().to_string() + } + + #[getter] + fn value<'py>(&self, py: Python<'py>) -> PyResult>> { + match self.metric.value() { + MetricValue::OutputRows(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())), + MetricValue::OutputBytes(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())), + MetricValue::ElapsedCompute(t) => Ok(Some(t.value().into_pyobject(py)?.into_any())), + MetricValue::SpillCount(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())), + MetricValue::SpilledBytes(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())), + MetricValue::SpilledRows(c) => Ok(Some(c.value().into_pyobject(py)?.into_any())), + MetricValue::CurrentMemoryUsage(g) => Ok(Some(g.value().into_pyobject(py)?.into_any())), + MetricValue::Count { count, .. } => { + Ok(Some(count.value().into_pyobject(py)?.into_any())) + } + MetricValue::Gauge { gauge, .. } => { + Ok(Some(gauge.value().into_pyobject(py)?.into_any())) + } + MetricValue::Time { time, .. } => Ok(Some(time.value().into_pyobject(py)?.into_any())), + MetricValue::StartTimestamp(ts) | MetricValue::EndTimestamp(ts) => { + Self::timestamp_to_pyobject(py, ts) + } + _ => Ok(None), + } + } + + #[getter] + fn value_as_datetime<'py>(&self, py: Python<'py>) -> PyResult>> { + match self.metric.value() { + MetricValue::StartTimestamp(ts) | MetricValue::EndTimestamp(ts) => { + Self::timestamp_to_pyobject(py, ts) + } + _ => Ok(None), + } + } + + #[getter] + fn partition(&self) -> Option { + self.metric.partition() + } + + fn labels(&self) -> HashMap { + self.metric + .labels() + .iter() + .map(|l| (l.name().to_string(), l.value().to_string())) + .collect() + } + + fn __repr__(&self) -> String { + format!("{}", self.metric.value()) + } +} diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index 8674a8b55..fac973884 100644 --- a/crates/core/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -26,6 +26,7 @@ use pyo3::types::PyBytes; use crate::context::PySessionContext; use crate::errors::PyDataFusionResult; +use crate::metrics::PyMetricsSet; #[pyclass( from_py_object, @@ -96,6 +97,10 @@ impl PyExecutionPlan { Ok(Self::new(plan)) } + pub fn metrics(&self) -> Option { + self.plan.metrics().map(PyMetricsSet::new) + } + fn __repr__(&self) -> String { self.display_indent() } diff --git a/docs/source/user-guide/dataframe/execution-metrics.rst b/docs/source/user-guide/dataframe/execution-metrics.rst new file mode 100644 index 000000000..764fa76ef --- /dev/null +++ b/docs/source/user-guide/dataframe/execution-metrics.rst @@ -0,0 +1,215 @@ +.. 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. + +.. _execution_metrics: + +Execution Metrics +================= + +Overview +-------- + +When DataFusion executes a query it compiles the logical plan into a tree of +*physical plan operators* (e.g. ``FilterExec``, ``ProjectionExec``, +``HashAggregateExec``). Each operator can record runtime statistics while it +runs. These statistics are called **execution metrics**. + +Typical metrics include: + +- **output_rows** – number of rows produced by the operator +- **elapsed_compute** – total CPU time (nanoseconds) spent inside the operator +- **spill_count** – number of times the operator spilled data to disk +- **spilled_bytes** – total bytes written to disk during spills +- **spilled_rows** – total rows written to disk during spills + +Metrics are collected *per-partition*: DataFusion may execute each operator +in parallel across several partitions. The convenience properties on +:py:class:`~datafusion.MetricsSet` (e.g. ``output_rows``, ``elapsed_compute``) +automatically sum the named metric across **all** partitions, giving a single +aggregate value for the operator as a whole. You can also access the raw +per-partition :py:class:`~datafusion.Metric` objects via +:py:meth:`~datafusion.MetricsSet.metrics`. + +When Are Metrics Available? +--------------------------- + +Some operators (for example ``DataSourceExec``) eagerly create a +:py:class:`~datafusion.MetricsSet` when the physical plan is built, so +:py:meth:`~datafusion.ExecutionPlan.metrics` may return a set even before any +rows have been processed. However, metric **values** such as ``output_rows`` +are only meaningful **after** the DataFrame has been executed via one of the +terminal operations: + +- :py:meth:`~datafusion.DataFrame.collect` +- :py:meth:`~datafusion.DataFrame.collect_partitioned` +- :py:meth:`~datafusion.DataFrame.execute_stream` + (metrics are available once the stream has been fully consumed) +- :py:meth:`~datafusion.DataFrame.execute_stream_partitioned` + (metrics are available once all partition streams have been fully consumed) + +Before execution, metric values will be ``0`` or ``None``. + +.. note:: + + **display() does not populate metrics.** + When a DataFrame is displayed in a notebook (e.g. via ``display(df)`` or + automatic ``repr`` output), DataFusion runs a *limited* internal execution + to fetch preview rows. This internal execution does **not** cache the + physical plan used, so :py:meth:`~datafusion.ExecutionPlan.collect_metrics` + will not reflect the display execution. To access metrics you must call + one of the terminal operations listed above. + +If you call :py:meth:`~datafusion.DataFrame.collect` (or another terminal +operation) multiple times on the same DataFrame, each call creates a fresh +physical plan. Metrics from :py:meth:`~datafusion.DataFrame.execution_plan` +always reflect the **most recent** execution. + +Reading the Physical Plan Tree +-------------------------------- + +:py:meth:`~datafusion.DataFrame.execution_plan` returns the root +:py:class:`~datafusion.ExecutionPlan` node of the physical plan tree. The tree +mirrors the operator pipeline: the root is typically a projection or +coalescing node; its children are filters, aggregates, scans, etc. + +The ``operator_name`` string returned by +:py:meth:`~datafusion.ExecutionPlan.collect_metrics` is the *display* name of +the node, for example ``"FilterExec: column1@0 > 1"``. This is the same string +you would see when calling ``plan.display()``. + +Aggregated vs Per-Partition Metrics +------------------------------------ + +DataFusion executes each operator across one or more **partitions** in +parallel. The :py:class:`~datafusion.MetricsSet` convenience properties +(``output_rows``, ``elapsed_compute``, etc.) automatically **sum** the named +metric across all partitions, giving a single aggregate value. + +To inspect individual partitions — for example to detect data skew where one +partition processes far more rows than others — iterate over the raw +:py:class:`~datafusion.Metric` objects: + +.. code-block:: python + + for metric in metrics_set.metrics(): + print(f" partition={metric.partition} {metric.name}={metric.value}") + +The ``partition`` property is a 0-based index (``0``, ``1``, …) identifying +which parallel slot processed this metric. It is ``None`` for metrics that +apply globally (not tied to a specific partition). + +Available Metrics +----------------- + +The following metrics are directly accessible as properties on +:py:class:`~datafusion.MetricsSet`: + +.. list-table:: + :header-rows: 1 + :widths: 25 75 + + * - Property + - Description + * - ``output_rows`` + - Number of rows emitted by the operator (summed across partitions). + * - ``elapsed_compute`` + - Wall-clock CPU time **in nanoseconds** spent inside the operator's + compute loop, excluding I/O wait. Useful for identifying which + operators are most expensive (summed across partitions). + * - ``spill_count`` + - Number of spill-to-disk events triggered by memory pressure. This is + a unitless count of events, not a measure of data volume (summed across + partitions). + * - ``spilled_bytes`` + - Total bytes written to disk during spill events (summed across + partitions). + * - ``spilled_rows`` + - Total rows written to disk during spill events (summed across + partitions). + +Any metric not listed above can be accessed via +:py:meth:`~datafusion.MetricsSet.sum_by_name`, or by iterating over the raw +:py:class:`~datafusion.Metric` objects returned by +:py:meth:`~datafusion.MetricsSet.metrics`. + +Labels +------ + +A :py:class:`~datafusion.Metric` may carry *labels*: key/value pairs that +provide additional context. Labels are operator-specific; most metrics have +an empty label dict. + +Some operators tag their metrics with labels to distinguish variants. For +example, a ``HashAggregateExec`` may record separate ``output_rows`` metrics +for intermediate and final output: + +.. code-block:: python + + for metric in metrics_set.metrics(): + print(metric.name, metric.labels()) + # output_rows {'output_type': 'final'} + # output_rows {'output_type': 'intermediate'} + +When summing by name (via :py:attr:`~datafusion.MetricsSet.output_rows` or +:py:meth:`~datafusion.MetricsSet.sum_by_name`), **all** metrics with that +name are summed regardless of labels. To filter by label, iterate over the +raw :py:class:`~datafusion.Metric` objects directly. + +End-to-End Example +------------------ + +.. code-block:: python + + from datafusion import SessionContext + + ctx = SessionContext() + ctx.sql("CREATE TABLE sales AS VALUES (1, 100), (2, 200), (3, 50)") + + df = ctx.sql("SELECT * FROM sales WHERE column1 > 1") + + # Execute the query — this populates the metrics + results = df.collect() + + # Retrieve the physical plan with metrics + plan = df.execution_plan() + + # Walk every operator and print its metrics + for operator_name, ms in plan.collect_metrics(): + if ms.output_rows is not None: + print(f"{operator_name}") + print(f" output_rows = {ms.output_rows}") + print(f" elapsed_compute = {ms.elapsed_compute} ns") + + # Access raw per-partition metrics + for operator_name, ms in plan.collect_metrics(): + for metric in ms.metrics(): + print( + f" partition={metric.partition} " + f"{metric.name}={metric.value} " + f"labels={metric.labels()}" + ) + +API Reference +------------- + +- :py:class:`datafusion.ExecutionPlan` — physical plan node +- :py:meth:`datafusion.ExecutionPlan.collect_metrics` — walk the tree and + return ``(operator_name, MetricsSet)`` pairs +- :py:meth:`datafusion.ExecutionPlan.metrics` — return the + :py:class:`~datafusion.MetricsSet` for a single node +- :py:class:`datafusion.MetricsSet` — aggregated metrics for one operator +- :py:class:`datafusion.Metric` — a single per-partition metric value diff --git a/docs/source/user-guide/dataframe/index.rst b/docs/source/user-guide/dataframe/index.rst index 510bcbc68..8475a7bd7 100644 --- a/docs/source/user-guide/dataframe/index.rst +++ b/docs/source/user-guide/dataframe/index.rst @@ -365,7 +365,16 @@ DataFusion provides many built-in functions for data manipulation: For a complete list of available functions, see the :py:mod:`datafusion.functions` module documentation. +Execution Metrics +----------------- + +After executing a DataFrame (via ``collect()``, ``execute_stream()``, etc.), +DataFusion populates per-operator runtime statistics such as row counts and +compute time. See :doc:`execution-metrics` for a full explanation and +worked example. + .. toctree:: :maxdepth: 1 rendering + execution-metrics diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index ee02c921d..80dfa2fab 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -56,7 +56,7 @@ from .expr import Expr, WindowFrame from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions -from .plan import ExecutionPlan, LogicalPlan +from .plan import ExecutionPlan, LogicalPlan, Metric, MetricsSet from .record_batch import RecordBatch, RecordBatchStream from .user_defined import ( Accumulator, @@ -86,6 +86,8 @@ "Expr", "InsertOp", "LogicalPlan", + "Metric", + "MetricsSet", "ParquetColumnOptions", "ParquetWriterOptions", "RecordBatch", diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 9c96a18fc..c0cfd523f 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -24,11 +24,15 @@ import datafusion._internal as df_internal if TYPE_CHECKING: + import datetime + from datafusion.context import SessionContext __all__ = [ "ExecutionPlan", "LogicalPlan", + "Metric", + "MetricsSet", ] @@ -151,3 +155,176 @@ def to_proto(self) -> bytes: Tables created in memory from record batches are currently not supported. """ return self._raw_plan.to_proto() + + def metrics(self) -> MetricsSet | None: + """Return metrics for this plan node, or None if this plan has no MetricsSet. + + Some operators (e.g. DataSourceExec) eagerly initialize a MetricsSet + when the plan is created, so this may return a set even before + execution. Metric *values* (such as ``output_rows``) are only + meaningful after the DataFrame has been executed. + """ + raw = self._raw_plan.metrics() + if raw is None: + return None + return MetricsSet(raw) + + def collect_metrics(self) -> list[tuple[str, MetricsSet]]: + """Return runtime statistics for each step of the query execution. + + DataFusion executes a query as a pipeline of operators — for example a + data source scan, followed by a filter, followed by a projection. After + the DataFrame has been executed (via + :py:meth:`~datafusion.DataFrame.collect`, + :py:meth:`~datafusion.DataFrame.execute_stream`, etc.), each operator + records statistics such as how many rows it produced and how much CPU + time it consumed. + + Each entry in the returned list corresponds to one operator that + recorded metrics. The first element of the tuple is the operator's + description string — the same text shown by + :py:meth:`display_indent` — which identifies both the operator type + and its key parameters, for example ``"FilterExec: column1@0 > 1"`` + or ``"DataSourceExec: partitions=1"``. + + Returns: + A list of ``(description, MetricsSet)`` tuples ordered from the + outermost operator (top of the execution tree) down to the + data-source leaves. Only operators that recorded at least one + metric are included. Returns an empty list if called before the + DataFrame has been executed. + """ + result: list[tuple[str, MetricsSet]] = [] + + def _walk(node: ExecutionPlan) -> None: + ms = node.metrics() + if ms is not None: + result.append((node.display(), ms)) + for child in node.children(): + _walk(child) + + _walk(self) + return result + + +class MetricsSet: + """A set of metrics for a single execution plan operator. + + A physical plan operator runs independently across one or more partitions. + :py:meth:`metrics` returns the raw per-partition :py:class:`Metric` objects. + The convenience properties (:py:attr:`output_rows`, :py:attr:`elapsed_compute`, + etc.) automatically sum the named metric across *all* partitions, giving a + single aggregate value for the operator as a whole. + """ + + def __init__(self, raw: df_internal.MetricsSet) -> None: + """This constructor should not be called by the end user.""" + self._raw = raw + + def metrics(self) -> list[Metric]: + """Return all individual metrics in this set.""" + return [Metric(m) for m in self._raw.metrics()] + + @property + def output_rows(self) -> int | None: + """Sum of output_rows across all partitions.""" + return self._raw.output_rows() + + @property + def elapsed_compute(self) -> int | None: + """Total CPU time (in nanoseconds) spent inside this operator's execute loop. + + Summed across all partitions. Returns ``None`` if no ``elapsed_compute`` + metric was recorded. + """ + return self._raw.elapsed_compute() + + @property + def spill_count(self) -> int | None: + """Number of times this operator spilled data to disk due to memory pressure. + + This is a count of spill events, not a byte count. Summed across all + partitions. Returns ``None`` if no ``spill_count`` metric was recorded. + """ + return self._raw.spill_count() + + @property + def spilled_bytes(self) -> int | None: + """Sum of spilled_bytes across all partitions.""" + return self._raw.spilled_bytes() + + @property + def spilled_rows(self) -> int | None: + """Sum of spilled_rows across all partitions.""" + return self._raw.spilled_rows() + + def sum_by_name(self, name: str) -> int | None: + """Sum the named metric across all partitions. + + Useful for accessing any metric not exposed as a first-class property. + Returns ``None`` if no metric with the given name was recorded. + + Args: + name: The metric name, e.g. ``"output_rows"`` or ``"elapsed_compute"``. + """ + return self._raw.sum_by_name(name) + + def __repr__(self) -> str: + """Return a string representation of the metrics set.""" + return repr(self._raw) + + +class Metric: + """A single execution metric with name, value, partition, and labels.""" + + def __init__(self, raw: df_internal.Metric) -> None: + """This constructor should not be called by the end user.""" + self._raw = raw + + @property + def name(self) -> str: + """The name of this metric (e.g. ``output_rows``).""" + return self._raw.name + + @property + def value(self) -> int | datetime.datetime | None: + """The value of this metric. + + Returns an ``int`` for counters, gauges, and time-based metrics + (nanoseconds), a :py:class:`~datetime.datetime` (UTC) for + ``start_timestamp`` / ``end_timestamp`` metrics, or ``None`` + when the value has not been set or is not representable. + """ + return self._raw.value + + @property + def value_as_datetime(self) -> datetime.datetime | None: + """The value as a UTC :py:class:`~datetime.datetime` for timestamp metrics. + + Returns ``None`` for all non-timestamp metrics and for timestamp + metrics whose value has not been set (e.g. before execution). + """ + return self._raw.value_as_datetime + + @property + def partition(self) -> int | None: + """The 0-based partition index this metric applies to. + + Returns ``None`` for metrics that are not partition-specific (i.e. they + apply globally across all partitions of the operator). + """ + return self._raw.partition + + def labels(self) -> dict[str, str]: + """Return the labels associated with this metric. + + Labels provide additional context for a metric. For example:: + + metric.labels() + # {'output_type': 'final'} + """ + return self._raw.labels() + + def __repr__(self) -> str: + """Return a string representation of the metric.""" + return repr(self._raw) diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 396acbe97..3705fc7ef 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -15,8 +15,16 @@ # specific language governing permissions and limitations # under the License. +import datetime + import pytest -from datafusion import ExecutionPlan, LogicalPlan, SessionContext +from datafusion import ( + ExecutionPlan, + LogicalPlan, + Metric, + MetricsSet, + SessionContext, +) # Note: We must use CSV because memory tables are currently not supported for @@ -40,3 +48,185 @@ def test_logical_plan_to_proto(ctx, df) -> None: execution_plan = ExecutionPlan.from_proto(ctx, execution_plan_bytes) assert str(original_execution_plan) == str(execution_plan) + + +def test_metrics_tree_walk() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + df.collect() + plan = df.execution_plan() + + results = plan.collect_metrics() + assert len(results) >= 1 + output_rows_by_op: dict[str, int] = {} + for name, ms in results: + assert isinstance(name, str) + assert isinstance(ms, MetricsSet) + if ms.output_rows is not None: + output_rows_by_op[name] = ms.output_rows + + # The filter passes rows where column1 > 1, so exactly + # 2 rows from (1,'a'),(2,'b'),(3,'c'). + # At least one operator must report exactly 2 output rows (the filter). + assert 2 in output_rows_by_op.values(), ( + f"Expected an operator with output_rows=2, got {output_rows_by_op}" + ) + + +def test_metric_properties() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + df.collect() + plan = df.execution_plan() + + found_any_metric = False + for _, ms in plan.collect_metrics(): + r = repr(ms) + assert isinstance(r, str) + for metric in ms.metrics(): + found_any_metric = True + assert isinstance(metric, Metric) + assert isinstance(metric.name, str) + assert len(metric.name) > 0 + assert metric.partition is None or isinstance(metric.partition, int) + assert metric.value is None or isinstance( + metric.value, int | datetime.datetime + ) + assert isinstance(metric.labels(), dict) + mr = repr(metric) + assert isinstance(mr, str) + assert len(mr) > 0 + assert found_any_metric, "Expected at least one metric after execution" + + +def test_no_meaningful_metrics_before_execution() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + plan_before = df.execution_plan() + + # Some plan nodes (e.g. DataSourceExec) eagerly initialize a MetricsSet, + # so metrics() may return a set even before execution. However, no rows + # should have been processed yet — output_rows must be absent or zero. + for _, ms in plan_before.collect_metrics(): + rows = ms.output_rows + assert rows is None or rows == 0, ( + f"Expected 0 output_rows before execution, got {rows}" + ) + + # After execution, at least one operator must report rows processed. + df.collect() + plan_after = df.execution_plan() + output_rows_after = [ + ms.output_rows + for _, ms in plan_after.collect_metrics() + if ms.output_rows is not None and ms.output_rows > 0 + ] + assert len(output_rows_after) > 0, "Expected output_rows > 0 after execution" + + +def test_collect_partitioned_metrics() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + + df.collect_partitioned() + plan = df.execution_plan() + + output_rows_values = [ + ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None + ] + assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}" + + +def test_execute_stream_metrics() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + + for _ in df.execute_stream(): + pass + + plan = df.execution_plan() + output_rows_values = [ + ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None + ] + assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}" + + +def test_execute_stream_partitioned_metrics() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + + for stream in df.execute_stream_partitioned(): + for _ in stream: + pass + + plan = df.execution_plan() + output_rows_values = [ + ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None + ] + assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}" + + +def test_value_as_datetime() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + df.collect() + plan = df.execution_plan() + + for _, ms in plan.collect_metrics(): + for metric in ms.metrics(): + if metric.name in ("start_timestamp", "end_timestamp"): + dt = metric.value_as_datetime + assert dt is None or isinstance(dt, datetime.datetime) + if dt is not None: + assert dt.tzinfo is not None + else: + assert metric.value_as_datetime is None + + +def test_metric_names_and_labels() -> None: + """Verify that known metric names appear and labels are well-formed.""" + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + df.collect() + plan = df.execution_plan() + + all_metric_names: set[str] = set() + for _, ms in plan.collect_metrics(): + for metric in ms.metrics(): + all_metric_names.add(metric.name) + # Labels must be a dict of str->str + labels = metric.labels() + for k, v in labels.items(): + assert isinstance(k, str) + assert isinstance(v, str) + + # After a filter query, we expect at minimum these standard metric names. + assert "output_rows" in all_metric_names, ( + f"Expected 'output_rows' in {all_metric_names}" + ) + assert "elapsed_compute" in all_metric_names, ( + f"Expected 'elapsed_compute' in {all_metric_names}" + ) + + +def test_collect_twice_has_metrics() -> None: + ctx = SessionContext() + ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") + df = ctx.sql("SELECT * FROM t WHERE column1 > 1") + + df.collect() + df.collect() + + plan = df.execution_plan() + output_rows_values = [ + ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None + ] + assert len(output_rows_values) > 0 From 398980d1edbb8ad6d9744236f2dfe0c6ab4b4665 Mon Sep 17 00:00:00 2001 From: Zeel Desai <72783325+zeel2104@users.noreply.github.com> Date: Mon, 13 Apr 2026 09:24:56 -0400 Subject: [PATCH 21/83] Support None comparisons for null expressions (#1489) * Support None comparisons for null expressions * Fold None comparison coverage into relational expr test --- python/datafusion/expr.py | 4 ++++ python/tests/test_expr.py | 8 ++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 7cd74ecd5..32004656f 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -483,6 +483,8 @@ def __eq__(self, rhs: object) -> Expr: Accepts either an expression or any valid PyArrow scalar literal value. """ + if rhs is None: + return self.is_null() if not isinstance(rhs, Expr): rhs = Expr.literal(rhs) return Expr(self.expr.__eq__(rhs.expr)) @@ -492,6 +494,8 @@ def __ne__(self, rhs: object) -> Expr: Accepts either an expression or any valid PyArrow scalar literal value. """ + if rhs is None: + return self.is_not_null() if not isinstance(rhs, Expr): rhs = Expr.literal(rhs) return Expr(self.expr.__ne__(rhs.expr)) diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index 1cf824a15..d046eb48c 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -153,8 +153,8 @@ def test_relational_expr(test_ctx): batch = pa.RecordBatch.from_arrays( [ - pa.array([1, 2, 3]), - pa.array(["alpha", "beta", "gamma"], type=pa.string_view()), + pa.array([1, 2, 3, None]), + pa.array(["alpha", "beta", "gamma", None], type=pa.string_view()), ], names=["a", "b"], ) @@ -171,6 +171,10 @@ def test_relational_expr(test_ctx): assert df.filter(col("b") != "beta").count() == 2 assert df.filter(col("a") == "beta").count() == 0 + assert df.filter(col("a") == None).count() == 1 # noqa: E711 + assert df.filter(col("a") != None).count() == 3 # noqa: E711 + assert df.filter(col("b") == None).count() == 1 # noqa: E711 + assert df.filter(col("b") != None).count() == 3 # noqa: E711 def test_expr_to_variant(): From 2715a32e939d17222c18e8adacf85ee45da464b9 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 14 Apr 2026 03:27:00 -0400 Subject: [PATCH 22/83] chore: update release documentation (#1494) * Update release documentation * Minor change to workflow because release start at 1 --- .../workflows/verify-release-candidate.yml | 2 +- dev/release/README.md | 60 ++++++++++--------- 2 files changed, 32 insertions(+), 30 deletions(-) diff --git a/.github/workflows/verify-release-candidate.yml b/.github/workflows/verify-release-candidate.yml index 7a5deff5b..6ecb547b5 100644 --- a/.github/workflows/verify-release-candidate.yml +++ b/.github/workflows/verify-release-candidate.yml @@ -27,7 +27,7 @@ on: required: true type: string rc_number: - description: Release candidate number (e.g., 0) + description: Release candidate number (e.g., 1) required: true type: string diff --git a/dev/release/README.md b/dev/release/README.md index ed28f4aa6..4833be55a 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -26,11 +26,11 @@ required due to changes in DataFusion rather than having a large amount of work is available. When there is a new official release of DataFusion, we update the `main` branch to point to that, update the version -number, and create a new release branch, such as `branch-0.8`. Once this branch is created, we switch the `main` branch +number, and create a new release branch, such as `branch-53`. Once this branch is created, we switch the `main` branch back to using GitHub dependencies. The release activity (such as generating the changelog) can then happen on the release branch without blocking ongoing development in the `main` branch. -We can cherry-pick commits from the `main` branch into `branch-0.8` as needed and then create new patch releases +We can cherry-pick commits from the `main` branch into `branch-53` as needed and then create new patch releases from that branch. ## Detailed Guide @@ -54,7 +54,8 @@ Before creating a new release: - We need to ensure that the main branch does not have any GitHub dependencies - a PR should be created and merged to update the major version number of the project -- A new release branch should be created, such as `branch-0.8` +- A new release branch should be created, such as `branch-53` +- It is best to push this branch to the apache repository rather than a personal fork in case patch releases are required. ## Preparing a Release Candidate @@ -65,14 +66,14 @@ We maintain a `CHANGELOG.md` so our users know what has been changed between rel The changelog is generated using a Python script: ```bash -$ GITHUB_TOKEN= ./dev/release/generate-changelog.py 24.0.0 HEAD 25.0.0 > dev/changelog/25.0.0.md +$ GITHUB_TOKEN= ./dev/release/generate-changelog.py 52.0.0 HEAD 53.0.0 > dev/changelog/53.0.0.md ``` This script creates a changelog from GitHub PRs based on the labels associated with them as well as looking for titles starting with `feat:`, `fix:`, or `docs:` . The script will produce output similar to: ``` -Fetching list of commits between 24.0.0 and HEAD +Fetching list of commits between 52.0.0 and HEAD Fetching pull requests Categorizing pull requests Generating changelog content @@ -81,6 +82,7 @@ Generating changelog content ### Update the version number The only place you should need to update the version is in the root `Cargo.toml`. +You will need to update this both in the workspace section and also in the dependencies. After updating the toml file, run `cargo update` to update the cargo lock file. If you do not want to update all the dependencies, you can instead run `cargo build` which should only update the version number for `datafusion-python`. @@ -94,14 +96,14 @@ you need to push a tag to start the CI process for release candidates. The follo the upstream repository is called `apache`. ```bash -git tag 0.8.0-rc1 -git push apache 0.8.0-rc1 +git tag 53.0.0-rc1 +git push apache 53.0.0-rc1 ``` ### Create a source release ```bash -./dev/release/create-tarball.sh 0.8.0 1 +./dev/release/create-tarball.sh 53.0.0 1 ``` This will also create the email template to send to the mailing list. @@ -124,10 +126,10 @@ Click on the action and scroll down to the bottom of the page titled "Artifacts" contain files such as: ```text -datafusion-22.0.0-cp37-abi3-macosx_10_7_x86_64.whl -datafusion-22.0.0-cp37-abi3-macosx_11_0_arm64.whl -datafusion-22.0.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl -datafusion-22.0.0-cp37-abi3-win_amd64.whl +datafusion-53.0.0-cp37-abi3-macosx_10_7_x86_64.whl +datafusion-53.0.0-cp37-abi3-macosx_11_0_arm64.whl +datafusion-53.0.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl +datafusion-53.0.0-cp37-abi3-win_amd64.whl ``` Upload the wheels to testpypi. @@ -135,23 +137,23 @@ Upload the wheels to testpypi. ```bash unzip dist.zip python3 -m pip install --upgrade setuptools twine build -python3 -m twine upload --repository testpypi datafusion-22.0.0-cp37-abi3-*.whl +python3 -m twine upload --repository testpypi datafusion-53.0.0-cp37-abi3-*.whl ``` When prompted for username, enter `__token__`. When prompted for a password, enter a valid GitHub Personal Access Token #### Publish Python Source Distribution to testpypi -Download the source tarball created in the previous step, untar it, and run: +Download the source tarball from the Apache server created in the previous step, untar it, and run: ```bash maturin sdist ``` -This will create a file named `dist/datafusion-0.7.0.tar.gz`. Upload this to testpypi: +This will create a file named `dist/datafusion-53.0.0.tar.gz`. Upload this to testpypi: ```bash -python3 -m twine upload --repository testpypi dist/datafusion-0.7.0.tar.gz +python3 -m twine upload --repository testpypi dist/datafusion-53.0.0.tar.gz ``` ### Run Verify Release Candidate Workflow @@ -162,8 +164,8 @@ Before sending the vote email, run the manually triggered GitHub Actions workflo 1. Go to https://github.com/apache/datafusion-python/actions/workflows/verify-release-candidate.yml 2. Click "Run workflow" -3. Set `version` to the release version (for example, `52.0.0`) -4. Set `rc_number` to the RC number (for example, `0`) +3. Set `version` to the release version (for example, `53.0.0`) +4. Set `rc_number` to the RC number (for example, `1`) 5. Wait for all jobs to complete successfully Include a short note in the vote email template that this workflow was run across all OS/architecture @@ -183,7 +185,7 @@ Releases may be verified using `verify-release-candidate.sh`: ```bash git clone https://github.com/apache/datafusion-python.git -dev/release/verify-release-candidate.sh 48.0.0 1 +dev/release/verify-release-candidate.sh 53.0.0 1 ``` Alternatively, one can run unit tests against a testpypi release candidate: @@ -195,7 +197,7 @@ cd datafusion-python # checkout the release commit git fetch --tags -git checkout 40.0.0-rc1 +git checkout 53.0.0-rc1 git submodule update --init --recursive # create the env @@ -203,7 +205,7 @@ python3 -m venv .venv source .venv/bin/activate # install release candidate -pip install --extra-index-url https://test.pypi.org/simple/ datafusion==40.0.0 +pip install --extra-index-url https://test.pypi.org/simple/ datafusion==53.0.0 # install test dependencies pip install pytest numpy pytest-asyncio @@ -224,7 +226,7 @@ Once the vote passes, we can publish the release. Create the source release tarball: ```bash -./dev/release/release-tarball.sh 0.8.0 1 +./dev/release/release-tarball.sh 53.0.0 1 ``` ### Publishing Rust Crate to crates.io @@ -232,7 +234,7 @@ Create the source release tarball: Some projects depend on the Rust crate directly, so we publish this to crates.io ```shell -cargo publish +cargo publish --workspace ``` ### Publishing Python Artifacts to PyPi @@ -252,15 +254,15 @@ Pypi packages auto upload to conda-forge via [datafusion feedstock](https://gith ### Push the Release Tag ```bash -git checkout 0.8.0-rc1 -git tag 0.8.0 -git push apache 0.8.0 +git checkout 53.0.0-rc1 +git tag 53.0.0 +git push apache 53.0.0 ``` ### Add the release to Apache Reporter Add the release to https://reporter.apache.org/addrelease.html?datafusion with a version name prefixed with `DATAFUSION-PYTHON`, -for example `DATAFUSION-PYTHON-31.0.0`. +for example `DATAFUSION-PYTHON-53.0.0`. The release information is used to generate a template for a board report (see example from Apache Arrow [here](https://github.com/apache/arrow/pull/14357)). @@ -283,7 +285,7 @@ svn ls https://dist.apache.org/repos/dist/dev/datafusion | grep datafusion-pytho Delete a release candidate: ```bash -svn delete -m "delete old DataFusion RC" https://dist.apache.org/repos/dist/dev/datafusion/apache-datafusion-python-7.1.0-rc1/ +svn delete -m "delete old DataFusion RC" https://dist.apache.org/repos/dist/dev/datafusion/apache-datafusion-python-53.0.0-rc1/ ``` #### Deleting old releases from `release` svn @@ -299,5 +301,5 @@ svn ls https://dist.apache.org/repos/dist/release/datafusion | grep datafusion-p Delete a release: ```bash -svn delete -m "delete old DataFusion release" https://dist.apache.org/repos/dist/release/datafusion/datafusion-python-7.0.0 +svn delete -m "delete old DataFusion release" https://dist.apache.org/repos/dist/release/datafusion/datafusion-python-52.0.0 ``` From 60d8b5dbb5e409cd9ce7692972420e955b8a802e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 14 Apr 2026 03:31:01 -0400 Subject: [PATCH 23/83] Fix error on show() with an explain plan (#1492) --- crates/core/src/dataframe.rs | 12 ++++++++++-- python/tests/test_dataframe.py | 10 ++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/crates/core/src/dataframe.rs b/crates/core/src/dataframe.rs index 2d815ec76..2e74991b8 100644 --- a/crates/core/src/dataframe.rs +++ b/crates/core/src/dataframe.rs @@ -38,8 +38,8 @@ use datafusion::dataframe::{DataFrame, DataFrameWriteOptions}; use datafusion::error::DataFusionError; use datafusion::execution::SendableRecordBatchStream; use datafusion::execution::context::TaskContext; -use datafusion::logical_expr::SortExpr; use datafusion::logical_expr::dml::InsertOp; +use datafusion::logical_expr::{LogicalPlan, SortExpr}; use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; use datafusion::physical_plan::{ ExecutionPlan as DFExecutionPlan, collect as df_collect, @@ -707,7 +707,15 @@ impl PyDataFrame { /// Print the result, 20 lines by default #[pyo3(signature = (num=20))] fn show(&self, py: Python, num: usize) -> PyDataFusionResult<()> { - let df = self.df.as_ref().clone().limit(0, Some(num))?; + let mut df = self.df.as_ref().clone(); + df = match self.df.logical_plan() { + LogicalPlan::Explain(_) | LogicalPlan::Analyze(_) => { + // Explain and Analyzer require they are at the top + // of the plan, so do not add a limit. + df + } + _ => df.limit(0, Some(num))?, + }; print_dataframe(py, df) } diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index bb8e9685c..091fa9b56 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -412,6 +412,16 @@ def test_show_empty(df, capsys): assert "DataFrame has no rows" in captured.out +def test_show_on_explain(ctx, capsys): + ctx.sql("explain select 1").show() + captured = capsys.readouterr() + assert "1 as Int64(1)" in captured.out + + ctx.sql("explain analyze select 1").show() + captured = capsys.readouterr() + assert "1 as Int64(1)" in captured.out + + def test_sort(df): df = df.sort(column("b").sort(ascending=False)) From 40309978c920bd123a4c7b764a2ddfdb97758607 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 23 Apr 2026 18:28:55 -0400 Subject: [PATCH 24/83] Add SKILL.md and enrich package docstring (#1497) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add AGENTS.md and enrich __init__.py module docstring Add python/datafusion/AGENTS.md as a comprehensive DataFrame API guide for AI agents and users. It ships with pip automatically (Maturin includes everything under python-source = "python"). Covers core abstractions, import conventions, data loading, all DataFrame operations, expression building, a SQL-to-DataFrame reference table, common pitfalls, idiomatic patterns, and a categorized function index. Enrich the __init__.py module docstring from 2 lines to a full overview with core abstractions, a quick-start example, and a pointer to AGENTS.md. Closes #1394 (PR 1a) Co-Authored-By: Claude Opus 4.6 (1M context) * Clarify audience of root vs package AGENTS.md The root AGENTS.md (symlinked as CLAUDE.md) is for contributors working on the project. Add a pointer to python/datafusion/AGENTS.md which is the user-facing DataFrame API guide shipped with the package. Also add the Apache license header to the package AGENTS.md. Co-Authored-By: Claude Opus 4.6 (1M context) * Add PR template and pre-commit check guidance to AGENTS.md Document that all PRs must follow .github/pull_request_template.md and that pre-commit hooks must pass before committing. List all configured hooks (actionlint, ruff, ruff-format, cargo fmt, cargo clippy, codespell, uv-lock) and the command to run them manually. Co-Authored-By: Claude Opus 4.6 (1M context) * Remove duplicated hook list from AGENTS.md Let the hooks be discoverable from .pre-commit-config.yaml rather than maintaining a separate list that can drift. Co-Authored-By: Claude Opus 4.6 (1M context) * Fix AGENTS.md: Arrow C Data Interface, aggregate filter, fluent example - Clarify that DataFusion works with any Arrow C Data Interface implementation, not just PyArrow. - Show the filter keyword argument on aggregate functions (the idiomatic HAVING equivalent) instead of the post-aggregate .filter() pattern. - Update the SQL reference table to show FILTER (WHERE ...) syntax. - Remove the now-incorrect "Aggregate then filter for HAVING" pitfall. - Add .collect() to the fluent chaining example so the result is clearly materialized. Co-Authored-By: Claude Opus 4.6 (1M context) * Update agents file after working through the first tpc-h query using only the text description * Add feedback from working through each of the TPC-H queries * Address Copilot review feedback on AGENTS.md - Wrap CASE/WHEN method-chain examples in parentheses and assign to a variable so they are valid Python as shown (Copilot #1, #2). - Fix INTERSECT/EXCEPT mapping: the default distinct=False corresponds to INTERSECT ALL / EXCEPT ALL, not the distinct forms. Updated both the Set Operations section and the SQL reference table to show both the ALL and distinct variants (Copilot #4). - Change write_parquet / write_csv / write_json examples to file-style paths (output.parquet, etc.) to match the convention used in existing tests and examples. Note that a directory path is also valid for partitioned output (Copilot #5). Verified INTERSECT/EXCEPT semantics with a script: df1.intersect(df2) -> [1, 1, 2] (= INTERSECT ALL) df1.intersect(df2, distinct=True) -> [1, 2] (= INTERSECT) Co-Authored-By: Claude Opus 4.6 (1M context) * Use short-form comparisons in AGENTS.md examples Drop lit() on the RHS of comparison operators since Expr auto-wraps raw Python values, matching the style the guide recommends (Copilot #3, #6). Updates examples in the Aggregation, CASE/WHEN, SQL reference table, Common Pitfalls, Fluent Chaining, and Variables-as-CTEs sections, plus the __init__.py quick-start snippet. Prose explanations of the rule (which cite the long form as the thing to avoid) are left unchanged. Co-Authored-By: Claude Opus 4.6 (1M context) * Move user guide from python/datafusion/AGENTS.md to SKILL.md The in-wheel AGENTS.md was not a real distribution channel -- no shipping agent walks site-packages for AGENTS.md files. Moving to SKILL.md at the repo root, with YAML frontmatter, lets the skill ecosystems (npx skills, Claude Code plugin marketplaces, community aggregators) discover it. Update the pointers in the contributor AGENTS.md and the __init__.py module docstring accordingly. The docstring now references the GitHub URL since the file no longer ships with the wheel. Co-Authored-By: Claude Opus 4.7 (1M context) * Address review feedback: doctest, streaming, date/timestamp - Convert the __init__.py quick-start block to doctest format so it is picked up by `pytest --doctest-modules` (already the project default), preventing silent rot. - Extract streaming into its own SKILL.md subsection with guidance on when to prefer execute_stream() over collect(), sync and async iteration, and execute_stream_partitioned() for per-partition streams. - Generalize the date-arithmetic rule from Date32 to both Date32 and Date64 (both reject Duration at any precision, both accept month_day_nano_interval), and note that Timestamp columns differ and do accept Duration. - Document the PyArrow-inherited type mapping returned by to_pydict()/to_pylist(), including the nanosecond fallback to pandas.Timestamp / pandas.Timedelta and the to_pandas() footgun where date columns come back as an object dtype. Co-Authored-By: Claude Opus 4.7 (1M context) * Distinguish user guide from agent reference in module docstring The docstring pointed readers at SKILL.md as a "comprehensive guide," but SKILL.md is written in a dense, skill-oriented format for agents — humans are better served by the online user guide. Put the online docs first as the primary reference and label the SKILL.md link as the agent reference. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- AGENTS.md | 34 +- SKILL.md | 733 ++++++++++++++++++++++++++++++++++ python/datafusion/__init__.py | 42 +- 3 files changed, 804 insertions(+), 5 deletions(-) create mode 100644 SKILL.md diff --git a/AGENTS.md b/AGENTS.md index 86c2e9c3b..7d3262710 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,7 +17,14 @@ under the License. --> -# Agent Instructions +# Agent Instructions for Contributors + +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). + +## Skills This project uses AI agent skills stored in `.ai/skills/`. Each skill is a directory containing a `SKILL.md` file with instructions for performing a specific task. @@ -26,6 +33,31 @@ 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. +## Pull Requests + +Every pull request must follow the template in +`.github/pull_request_template.md`. The description must include these sections: + +1. **Which issue does this PR close?** — Link the issue with `Closes #NNN`. +2. **Rationale for this change** — Why the change is needed (skip if the issue + already explains it clearly). +3. **What changes are included in this PR?** — Summarize the individual changes. +4. **Are there any user-facing changes?** — Note any changes visible to users + (new APIs, changed behavior, new files shipped in the package, etc.). If + there are breaking changes to public APIs, add the `api change` label. + +## Pre-commit Checks + +Always run pre-commit checks **before** committing. The hooks are defined in +`.pre-commit-config.yaml` and run automatically on `git commit` if pre-commit +is installed as a git hook. To run all hooks manually: + +```bash +pre-commit run --all-files +``` + +Fix any failures before committing. + ## Python Function Docstrings Every Python function must include a docstring with usage examples. diff --git a/SKILL.md b/SKILL.md new file mode 100644 index 000000000..9ba1c0cac --- /dev/null +++ b/SKILL.md @@ -0,0 +1,733 @@ + + +--- +name: datafusion-python +description: Use when the user is writing datafusion-python (Apache DataFusion Python bindings) DataFrame or SQL code. Covers imports, data loading, DataFrame operations, expression building, SQL-to-DataFrame mappings, idiomatic patterns, and common pitfalls. +--- + +# DataFusion Python DataFrame API Guide + +## What Is DataFusion? + +DataFusion is an **in-process query engine** built on Apache Arrow. It is not a +database -- there is no server, no connection string, and no external +dependencies. You create a `SessionContext`, point it at data (Parquet, CSV, +JSON, Arrow IPC, Pandas, Polars, or raw Python dicts/lists), and run queries +using either SQL or the DataFrame API described below. + +All data flows through **Apache Arrow**. The canonical Python implementation is +PyArrow (`pyarrow.RecordBatch` / `pyarrow.Table`), but any library that +conforms to the [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html) +can interoperate with DataFusion. + +## Core Abstractions + +| Abstraction | Role | Key import | +|---|---|---| +| `SessionContext` | Entry point. Loads data, runs SQL, produces DataFrames. | `from datafusion import SessionContext` | +| `DataFrame` | Lazy query builder. Each method returns a new DataFrame. | Returned by context methods | +| `Expr` | Expression tree node (column ref, literal, function call, ...). | `from datafusion import col, lit` | +| `functions` | 290+ built-in scalar, aggregate, and window functions. | `from datafusion import functions as F` | + +## Import Conventions + +```python +from datafusion import SessionContext, col, lit +from datafusion import functions as F +``` + +## Data Loading + +```python +ctx = SessionContext() + +# From files +df = ctx.read_parquet("path/to/data.parquet") +df = ctx.read_csv("path/to/data.csv") +df = ctx.read_json("path/to/data.json") + +# From Python objects +df = ctx.from_pydict({"a": [1, 2, 3], "b": ["x", "y", "z"]}) +df = ctx.from_pylist([{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]) +df = ctx.from_pandas(pandas_df) +df = ctx.from_polars(polars_df) +df = ctx.from_arrow(arrow_table) + +# From SQL +df = ctx.sql("SELECT a, b FROM my_table WHERE a > 1") +``` + +To make a DataFrame queryable by name in SQL, register it first: + +```python +ctx.register_parquet("my_table", "path/to/data.parquet") +ctx.register_csv("my_table", "path/to/data.csv") +``` + +## DataFrame Operations Quick Reference + +Every method returns a **new** DataFrame (immutable/lazy). Chain them fluently. + +### Projection + +```python +df.select("a", "b") # preferred: plain names as strings +df.select(col("a"), (col("b") + 1).alias("b_plus_1")) # use col()/Expr only when you need an expression + +df.with_column("new_col", col("a") + lit(10)) # add one column +df.with_columns( + col("a").alias("x"), + y=col("b") + lit(1), # named keyword form +) + +df.drop("unwanted_col") +df.with_column_renamed("old_name", "new_name") +``` + +When a column is referenced by name alone, pass the name as a string rather +than wrapping it in `col()`. Reach for `col()` only when the projection needs +arithmetic, aliasing, casting, or another expression operation. + +**Case sensitivity**: both `select("Name")` and `col("Name")` lowercase the +identifier. For a column whose real name has uppercase letters, embed double +quotes inside the string: `select('"MyCol"')` or `col('"MyCol"')`. Without the +inner quotes the lookup will fail with `No field named mycol`. + +### Filtering + +```python +df.filter(col("a") > 10) +df.filter(col("a") > 10, col("b") == "x") # multiple = AND +df.filter("a > 10") # SQL expression string +``` + +Raw Python values on the right-hand side of a comparison are auto-wrapped +into literals by the `Expr` operators, so prefer `col("a") > 10` over +`col("a") > lit(10)`. See the Comparisons section and pitfall #2 for the +full rule. + +### Aggregation + +```python +# GROUP BY a, compute sum(b) and count(*) +df.aggregate(["a"], [F.sum(col("b")), F.count(col("a"))]) + +# HAVING equivalent: use the filter keyword on the aggregate function +df.aggregate( + ["region"], + [F.sum(col("sales"), filter=col("sales") > 1000).alias("large_sales")], +) +``` + +As with `select()`, group keys can be passed as plain name strings. Reach for +`col(...)` only when the grouping expression needs arithmetic, aliasing, +casting, or another expression operation. + +Most aggregate functions accept an optional `filter` keyword argument. When +provided, only rows where the filter expression is true contribute to the +aggregate. + +### Sorting + +```python +df.sort(col("a")) # ascending (default) +df.sort(col("a").sort(ascending=False)) # descending +df.sort(col("a").sort(nulls_first=False)) # override null placement +``` + +A plain expression passed to `sort()` is already treated as ascending. Only +reach for `col(...).sort(...)` when you need to override a default (descending +order or null placement). Writing `col("a").sort(ascending=True)` is redundant. + +### Joining + +```python +# Equi-join on shared column name +df1.join(df2, on="key") +df1.join(df2, on="key", how="left") + +# Different column names +df1.join(df2, left_on="id", right_on="fk_id", how="inner") + +# Expression-based join (supports inequality predicates) +df1.join_on(df2, col("a") == col("b"), how="inner") + +# Semi join: keep rows from left where a match exists in right (like EXISTS) +df1.join(df2, on="key", how="semi") + +# Anti join: keep rows from left where NO match exists in right (like NOT EXISTS) +df1.join(df2, on="key", how="anti") +``` + +Join types: `"inner"`, `"left"`, `"right"`, `"full"`, `"semi"`, `"anti"`. + +Inner is the default `how`. Prefer `df1.join(df2, on="key")` over +`df1.join(df2, on="key", how="inner")` — drop `how=` unless you need a +non-inner join type. + +When the two sides' join columns have different native names, use +`left_on=`/`right_on=` with the original names rather than aliasing one side +to match the other — see pitfall #7. + +### Window Functions + +```python +from datafusion import WindowFrame + +# Row number partitioned by group, ordered by value +df.window( + F.row_number( + partition_by=[col("group")], + order_by=[col("value")], + ).alias("rn") +) + +# Using a Window object for reuse +from datafusion.expr import Window + +win = Window( + partition_by=[col("group")], + order_by=[col("value").sort(ascending=True)], +) +df.select( + col("group"), + col("value"), + F.sum(col("value")).over(win).alias("running_total"), +) + +# With explicit frame bounds +win = Window( + partition_by=[col("group")], + order_by=[col("value").sort(ascending=True)], + window_frame=WindowFrame("rows", 0, None), # current row to unbounded following +) +``` + +### Set Operations + +```python +df1.union(df2) # UNION ALL (by position) +df1.union(df2, distinct=True) # UNION DISTINCT +df1.union_by_name(df2) # match columns by name, not position +df1.intersect(df2) # INTERSECT ALL +df1.intersect(df2, distinct=True) # INTERSECT (distinct) +df1.except_all(df2) # EXCEPT ALL +df1.except_all(df2, distinct=True) # EXCEPT (distinct) +``` + +### Limit and Offset + +```python +df.limit(10) # first 10 rows +df.limit(10, offset=20) # skip 20, then take 10 +``` + +### Deduplication + +```python +df.distinct() # remove duplicate rows +df.distinct_on( # keep first row per group (like DISTINCT ON in Postgres) + [col("a")], # uniqueness columns + [col("a"), col("b")], # output columns + [col("b").sort(ascending=True)], # which row to keep +) +``` + +## Executing and Collecting Results + +DataFrames are lazy until you collect. + +```python +df.show() # print formatted table to stdout +batches = df.collect() # list[pa.RecordBatch] +arr = df.collect_column("col_name") # pa.Array | pa.ChunkedArray (single column) +table = df.to_arrow_table() # pa.Table +pandas_df = df.to_pandas() # pd.DataFrame +polars_df = df.to_polars() # pl.DataFrame +py_dict = df.to_pydict() # dict[str, list] +py_list = df.to_pylist() # list[dict] +count = df.count() # int +``` + +### Date and Timestamp Type Conversion + +The Python type returned by `to_pydict()` / `to_pylist()` depends on the Arrow +column type, and the mapping is inherited from PyArrow: + +| Arrow type | Python type returned | +|---|---| +| `timestamp(s)` / `(ms)` / `(us)` | `datetime.datetime` | +| `timestamp(ns)` | `pandas.Timestamp` | +| `date32` / `date64` | `datetime.date` | +| `duration(s)` / `(ms)` / `(us)` | `datetime.timedelta` | +| `duration(ns)` | `pandas.Timedelta` | + +The nanosecond-precision fallback to pandas types is the main surprise: +pandas is not a hard dependency of `datafusion`, but PyArrow reaches for it +when `datetime.datetime` / `datetime.timedelta` would lose precision (stdlib +types only go to microseconds). If you need plain stdlib types, cast to a +coarser unit before collecting, e.g. +`df.select(col("ts").cast(pa.timestamp("us")))`. + +`df.to_pandas()` has its own footgun for dates: pandas has no pure-date dtype, +so a `date32`/`date64` column comes back as an `object` column of +`datetime.date` values rather than `datetime64[ns]`. If downstream code +expects a datetime column, cast on the DataFusion side first: +`col("ship_date").cast(pa.timestamp("ns"))`. + +### Streaming Results + +Prefer streaming over `collect()` when the result is too large to materialize +in memory, when you want to start processing before the query finishes, or +when you may break out of the loop early. `execute_stream()` pulls one +`RecordBatch` at a time from the execution plan rather than buffering the +whole result up front. + +```python +# Single-partition stream; batch is a datafusion.RecordBatch +stream = df.execute_stream() +for batch in stream: + process(batch.to_pyarrow()) # convert to pa.RecordBatch if needed + +# DataFrame is iterable directly (delegates to execute_stream) +for batch in df: + process(batch.to_pyarrow()) + +# One stream per partition, for parallel consumption +for stream in df.execute_stream_partitioned(): + for batch in stream: + process(batch.to_pyarrow()) +``` + +Async iteration is also supported via `async for batch in df: ...` (or +`df.execute_stream()`), which is useful when batches are interleaved with +other I/O. + +### Writing Results + +```python +df.write_parquet("output.parquet") +df.write_csv("output.csv") +df.write_json("output.json") +``` + +You can also pass a directory path (e.g., `"output/"`) to write a multi-file +partitioned output. + +## Expression Building + +### Column References and Literals + +```python +col("column_name") # reference a column +lit(42) # integer literal +lit("hello") # string literal +lit(3.14) # float literal +lit(pa.scalar(value)) # PyArrow scalar (preserves Arrow type) +``` + +`lit()` accepts PyArrow scalars directly -- prefer this over converting Arrow +data to Python and back when working with values extracted from query results. + +### Arithmetic + +```python +col("price") * col("quantity") # multiplication +col("a") + lit(1) # addition +col("a") - col("b") # subtraction +col("a") / lit(2) # division +col("a") % lit(3) # modulo +``` + +### Date Arithmetic + +`Date32` and `Date64` columns both require `Interval` types for arithmetic, +not `Duration`. Use PyArrow's `month_day_nano_interval` type, which takes a +`(months, days, nanos)` tuple: + +```python +import pyarrow as pa + +# Subtract 90 days from a date column +col("ship_date") - lit(pa.scalar((0, 90, 0), type=pa.month_day_nano_interval())) + +# Subtract 3 months +col("ship_date") - lit(pa.scalar((3, 0, 0), type=pa.month_day_nano_interval())) +``` + +**Important**: `lit(datetime.timedelta(days=90))` creates a `Duration(µs)` +literal, which is **not** compatible with `Date32`/`Date64` arithmetic +(`Duration(ms)` and `Duration(ns)` are rejected too). Always use +`pa.month_day_nano_interval()` for date operations. + +**Timestamps behave differently**: `Timestamp` columns *do* accept `Duration`, +so `col("ts") - lit(datetime.timedelta(days=1))` works. The interval-only +rule applies specifically to date columns. + +### Comparisons + +```python +col("a") > 10 +col("a") >= 10 +col("a") < 10 +col("a") <= 10 +col("a") == "x" +col("a") != "x" +col("a") == None # same as col("a").is_null() +col("a") != None # same as col("a").is_not_null() +``` + +Comparison operators auto-wrap the right-hand Python value into a literal, +so writing `col("a") > lit(10)` is redundant. Drop the `lit()` in +comparisons. Reach for `lit()` only when auto-wrapping does not apply — see +pitfall #2. + +### Boolean Logic + +**Important**: Python's `and`, `or`, `not` keywords do NOT work with Expr +objects. You must use the bitwise operators: + +```python +(col("a") > 1) & (col("b") < 10) # AND +(col("a") > 1) | (col("b") < 10) # OR +~(col("a") > 1) # NOT +``` + +Always wrap each comparison in parentheses when combining with `&`, `|`, `~` +because Python's operator precedence for bitwise operators is different from +logical operators. + +### Null Handling + +```python +col("a").is_null() +col("a").is_not_null() +col("a").fill_null(lit(0)) # replace NULL with a value +F.coalesce(col("a"), col("b")) # first non-null value +F.nullif(col("a"), lit(0)) # return NULL if a == 0 +``` + +### CASE / WHEN + +```python +# Simple CASE (matching on a single expression) +status_label = ( + F.case(col("status")) + .when(lit("A"), lit("Active")) + .when(lit("I"), lit("Inactive")) + .otherwise(lit("Unknown")) +) + +# Searched CASE (each branch has its own predicate) +severity = ( + F.when(col("value") > 100, lit("high")) + .when(col("value") > 50, lit("medium")) + .otherwise(lit("low")) +) +``` + +### Casting + +```python +import pyarrow as pa + +col("a").cast(pa.float64()) +col("a").cast(pa.utf8()) +col("a").cast(pa.date32()) +``` + +### Aliasing + +```python +(col("a") + col("b")).alias("total") +``` + +### BETWEEN and IN + +```python +col("a").between(lit(1), lit(10)) # 1 <= a <= 10 +F.in_list(col("a"), [lit(1), lit(2), lit(3)]) # a IN (1, 2, 3) +F.in_list(col("a"), [lit(1), lit(2)], negated=True) # a NOT IN (1, 2) +``` + +### Struct and Array Access + +```python +col("struct_col")["field_name"] # access struct field +col("array_col")[0] # access array element (0-indexed) +col("array_col")[1:3] # array slice (0-indexed) +``` + +## SQL-to-DataFrame Reference + +| SQL | DataFrame API | +|---|---| +| `SELECT a, b` | `df.select("a", "b")` | +| `SELECT a, b + 1 AS c` | `df.select(col("a"), (col("b") + lit(1)).alias("c"))` | +| `SELECT *, a + 1 AS c` | `df.with_column("c", col("a") + lit(1))` | +| `WHERE a > 10` | `df.filter(col("a") > 10)` | +| `GROUP BY a` with `SUM(b)` | `df.aggregate(["a"], [F.sum(col("b"))])` | +| `SUM(b) FILTER (WHERE b > 100)` | `F.sum(col("b"), filter=col("b") > 100)` | +| `ORDER BY a DESC` | `df.sort(col("a").sort(ascending=False))` | +| `LIMIT 10 OFFSET 5` | `df.limit(10, offset=5)` | +| `DISTINCT` | `df.distinct()` | +| `a INNER JOIN b ON a.id = b.id` | `a.join(b, on="id")` | +| `a LEFT JOIN b ON a.id = b.fk` | `a.join(b, left_on="id", right_on="fk", how="left")` | +| `WHERE EXISTS (SELECT ...)` | `a.join(b, on="key", how="semi")` | +| `WHERE NOT EXISTS (SELECT ...)` | `a.join(b, on="key", how="anti")` | +| `UNION ALL` | `df1.union(df2)` | +| `UNION` (distinct) | `df1.union(df2, distinct=True)` | +| `INTERSECT ALL` | `df1.intersect(df2)` | +| `INTERSECT` (distinct) | `df1.intersect(df2, distinct=True)` | +| `EXCEPT ALL` | `df1.except_all(df2)` | +| `EXCEPT` (distinct) | `df1.except_all(df2, distinct=True)` | +| `CASE x WHEN 1 THEN 'a' END` | `F.case(col("x")).when(lit(1), lit("a")).end()` | +| `CASE WHEN x > 1 THEN 'a' END` | `F.when(col("x") > 1, lit("a")).end()` | +| `x IN (1, 2, 3)` | `F.in_list(col("x"), [lit(1), lit(2), lit(3)])` | +| `x BETWEEN 1 AND 10` | `col("x").between(lit(1), lit(10))` | +| `CAST(x AS DOUBLE)` | `col("x").cast(pa.float64())` | +| `ROW_NUMBER() OVER (...)` | `F.row_number(partition_by=[...], order_by=[...])` | +| `SUM(x) OVER (...)` | `F.sum(col("x")).over(window)` | +| `x IS NULL` | `col("x").is_null()` | +| `COALESCE(a, b)` | `F.coalesce(col("a"), col("b"))` | + +## Common Pitfalls + +1. **Boolean operators**: Use `&`, `|`, `~` -- not Python's `and`, `or`, `not`. + Always parenthesize: `(col("a") > 1) & (col("b") < 2)`. + +2. **Wrapping scalars with `lit()`**: Prefer raw Python values on the + right-hand side of comparisons — `col("a") > 10`, `col("name") == "Alice"` + — because the Expr comparison operators auto-wrap them. Writing + `col("a") > lit(10)` is redundant. Reserve `lit()` for places where + auto-wrapping does *not* apply: + - standalone scalars passed into function calls: + `F.coalesce(col("a"), lit(0))`, not `F.coalesce(col("a"), 0)` + - arithmetic between two literals with no column involved: + `lit(1) - col("discount")` is fine, but `lit(1) - lit(2)` needs both + - values that must carry a specific Arrow type, via `lit(pa.scalar(...))` + - `.when(...)`, `.otherwise(...)`, `F.nullif(...)`, `.between(...)`, + `F.in_list(...)` and similar method/function arguments + +3. **Column name quoting**: Column names are normalized to lowercase by default + in both `select("...")` and `col("...")`. To reference a column with + uppercase letters, use double quotes inside the string: + `select('"MyColumn"')` or `col('"MyColumn"')`. + +4. **DataFrames are immutable**: Every method returns a **new** DataFrame. You + must capture the return value: + ```python + df = df.filter(col("a") > 1) # correct + df.filter(col("a") > 1) # WRONG -- result is discarded + ``` + +5. **Window frame defaults**: When using `order_by` in a window, the default + frame is `RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`. For a full + partition frame, set `window_frame=WindowFrame("rows", None, None)`. + +6. **Arithmetic on aggregates belongs in a later `select`, not inside + `aggregate`**: Each item in the aggregate list must be a single aggregate + call (optionally aliased). Combining aggregates with arithmetic inside + `aggregate(...)` fails with `Internal error: Invalid aggregate expression`. + Alias the aggregates, then compute the combination downstream: + ```python + # WRONG -- arithmetic wraps two aggregates + df.aggregate([], [(lit(100) * F.sum(col("a")) / F.sum(col("b"))).alias("ratio")]) + + # CORRECT -- aggregate first, then combine + (df.aggregate([], [F.sum(col("a")).alias("num"), F.sum(col("b")).alias("den")]) + .select((lit(100) * col("num") / col("den")).alias("ratio"))) + ``` + +7. **Don't alias a join column to match the other side**: When equi-joining + with `on="key"`, renaming the join column on one side via `.alias("key")` + in a fresh projection creates a schema where one side's `key` is + qualified (`?table?.key`) and the other is unqualified. The join then + fails with `Schema contains qualified field name ... and unqualified + field name ... which would be ambiguous`. Use `left_on=`/`right_on=` with + the native names, or use `join_on(...)` with an explicit equality. + ```python + # WRONG -- alias on one side produces ambiguous schema after join + failed = orders.select(col("o_orderkey").alias("l_orderkey")) + li.join(failed, on="l_orderkey") # ambiguous l_orderkey error + + # CORRECT -- keep native names, use left_on/right_on + failed = orders.select("o_orderkey") + li.join(failed, left_on="l_orderkey", right_on="o_orderkey") + + # ALSO CORRECT -- explicit predicate via join_on + # (note: join_on keeps both key columns in the output, unlike on="key") + li.join_on(failed, col("l_orderkey") == col("o_orderkey")) + ``` + +## Idiomatic Patterns + +### Fluent Chaining + +```python +result = ( + ctx.read_parquet("data.parquet") + .filter(col("year") >= 2020) + .select(col("region"), col("sales")) + .aggregate(["region"], [F.sum(col("sales")).alias("total")]) + .sort(col("total").sort(ascending=False)) + .limit(10) +) +result.show() +``` + +### Using Variables as CTEs + +Instead of SQL CTEs (`WITH ... AS`), assign intermediate DataFrames to +variables: + +```python +base = ctx.read_parquet("orders.parquet").filter(col("status") == "shipped") +by_region = base.aggregate(["region"], [F.sum(col("amount")).alias("total")]) +top_regions = by_region.filter(col("total") > 10000) +``` + +### Reusing Expressions as Variables + +Just like DataFrames, expressions (`Expr`) can be stored in variables and used +anywhere an `Expr` is expected. This is useful for building up complex +expressions or reusing a computed value across multiple operations: + +```python +# Build an expression and reuse it +disc_price = col("price") * (lit(1) - col("discount")) +df = df.select( + col("id"), + disc_price.alias("disc_price"), + (disc_price * (lit(1) + col("tax"))).alias("total"), +) + +# Use a collected scalar as an expression +max_val = result_df.collect_column("max_price")[0] # PyArrow scalar +cutoff = lit(max_val) - lit(pa.scalar((0, 90, 0), type=pa.month_day_nano_interval())) +df = df.filter(col("ship_date") <= cutoff) # cutoff is already an Expr +``` + +**Important**: Do not wrap an `Expr` in `lit()`. `lit()` is for converting +Python/PyArrow values into expressions. If a value is already an `Expr`, use it +directly. + +### Window Functions for Scalar Subqueries + +Where SQL uses a correlated scalar subquery, the idiomatic DataFrame approach +is a window function: + +```sql +-- SQL scalar subquery +SELECT *, (SELECT SUM(b) FROM t WHERE t.group = s.group) AS group_total FROM s +``` + +```python +# DataFrame: window function +win = Window(partition_by=[col("group")]) +df = df.with_column("group_total", F.sum(col("b")).over(win)) +``` + +### Semi/Anti Joins for EXISTS / NOT EXISTS + +```sql +-- SQL: WHERE EXISTS (SELECT 1 FROM other WHERE other.key = main.key) +-- DataFrame: +result = main.join(other, on="key", how="semi") + +-- SQL: WHERE NOT EXISTS (SELECT 1 FROM other WHERE other.key = main.key) +-- DataFrame: +result = main.join(other, on="key", how="anti") +``` + +### Computed Columns + +```python +# Add computed columns while keeping all originals +df = df.with_column("full_name", F.concat(col("first"), lit(" "), col("last"))) +df = df.with_column("discounted", col("price") * lit(0.9)) +``` + +## Available Functions (Categorized) + +The `functions` module (imported as `F`) provides 290+ functions. Key categories: + +**Aggregate**: `sum`, `avg`, `min`, `max`, `count`, `count_star`, `median`, +`stddev`, `stddev_pop`, `var_samp`, `var_pop`, `corr`, `covar`, `approx_distinct`, +`approx_median`, `approx_percentile_cont`, `array_agg`, `string_agg`, +`first_value`, `last_value`, `bit_and`, `bit_or`, `bit_xor`, `bool_and`, +`bool_or`, `grouping`, `regr_*` (9 regression functions) + +**Window**: `row_number`, `rank`, `dense_rank`, `percent_rank`, `cume_dist`, +`ntile`, `lag`, `lead`, `first_value`, `last_value`, `nth_value` + +**String**: `length`, `lower`, `upper`, `trim`, `ltrim`, `rtrim`, `lpad`, +`rpad`, `starts_with`, `ends_with`, `contains`, `substr`, `substring`, +`replace`, `reverse`, `repeat`, `split_part`, `concat`, `concat_ws`, +`initcap`, `ascii`, `chr`, `left`, `right`, `strpos`, `translate`, `overlay`, +`levenshtein` + +`F.substr(str, start)` takes **only two arguments** and returns the tail of +the string from `start` onward — passing a third length argument raises +`TypeError: substr() takes 2 positional arguments but 3 were given`. For the +SQL-style 3-arg form (`SUBSTRING(str FROM start FOR length)`), use +`F.substring(col("s"), lit(start), lit(length))`. For a fixed-length prefix, +`F.left(col("s"), lit(n))` is cleanest. + +```python +# WRONG — substr does not accept a length argument +F.substr(col("c_phone"), lit(1), lit(2)) +# CORRECT +F.substring(col("c_phone"), lit(1), lit(2)) # explicit length +F.left(col("c_phone"), lit(2)) # prefix shortcut +``` + +**Math**: `abs`, `ceil`, `floor`, `round`, `trunc`, `sqrt`, `cbrt`, `exp`, +`ln`, `log`, `log2`, `log10`, `pow`, `signum`, `pi`, `random`, `factorial`, +`gcd`, `lcm`, `greatest`, `least`, sin/cos/tan and inverse/hyperbolic variants + +**Date/Time**: `now`, `today`, `current_date`, `current_time`, +`current_timestamp`, `date_part`, `date_trunc`, `date_bin`, `extract`, +`to_timestamp`, `to_timestamp_millis`, `to_timestamp_micros`, +`to_timestamp_nanos`, `to_timestamp_seconds`, `to_unixtime`, `from_unixtime`, +`make_date`, `make_time`, `to_date`, `to_time`, `to_local_time`, `date_format` + +**Conditional**: `case`, `when`, `coalesce`, `nullif`, `ifnull`, `nvl`, `nvl2` + +**Array/List**: `array`, `make_array`, `array_agg`, `array_length`, +`array_element`, `array_slice`, `array_append`, `array_prepend`, +`array_concat`, `array_has`, `array_has_all`, `array_has_any`, `array_position`, +`array_remove`, `array_distinct`, `array_sort`, `array_reverse`, `flatten`, +`array_to_string`, `array_intersect`, `array_union`, `array_except`, +`generate_series` +(Most `array_*` functions also have `list_*` aliases.) + +**Struct/Map**: `struct`, `named_struct`, `get_field`, `make_map`, `map_keys`, +`map_values`, `map_entries`, `map_extract` + +**Regex**: `regexp_like`, `regexp_match`, `regexp_replace`, `regexp_count`, +`regexp_instr` + +**Hash**: `md5`, `sha224`, `sha256`, `sha384`, `sha512`, `digest` + +**Type**: `arrow_typeof`, `arrow_cast`, `arrow_metadata` + +**Other**: `in_list`, `order_by`, `alias`, `col`, `encode`, `decode`, +`to_hex`, `to_char`, `uuid`, `version`, `bit_length`, `octet_length` diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 80dfa2fab..e4972411a 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -15,10 +15,44 @@ # specific language governing permissions and limitations # under the License. -"""DataFusion python package. - -This is a Python library that binds to Apache Arrow in-memory query engine DataFusion. -See https://datafusion.apache.org/python for more information. +"""DataFusion: an in-process query engine built on Apache Arrow. + +DataFusion is not a database -- it has no server and no external dependencies. +You create a :py:class:`SessionContext`, point it at data sources (Parquet, CSV, +JSON, Arrow IPC, Pandas, Polars, or raw Python dicts/lists), and run queries +using either SQL or the DataFrame API. + +Core abstractions +----------------- +- **SessionContext** -- entry point for loading data, running SQL, and creating + DataFrames. +- **DataFrame** -- lazy query builder. Every method returns a new DataFrame; + call :py:meth:`~datafusion.dataframe.DataFrame.collect` or a ``to_*`` + method to execute. +- **Expr** -- expression tree node for column references, literals, and function + calls. Build with :py:func:`col` and :py:func:`lit`. +- **functions** -- 290+ built-in scalar, aggregate, and window functions. + +Quick start +----------- + +>>> from datafusion import SessionContext, col +>>> from datafusion import functions as F +>>> ctx = SessionContext() +>>> df = ctx.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}) +>>> result = ( +... df.filter(col("a") > 1) +... .with_column("total", col("a") + col("b")) +... .aggregate([], [F.sum(col("total")).alias("grand_total")]) +... ) +>>> result.to_pydict() +{'grand_total': [16]} + +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 """ from __future__ import annotations From 8a5d783c7e418bfbbd95e48a2d9cacafea6162c7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 23 Apr 2026 19:05:06 -0400 Subject: [PATCH 25/83] Skills require the header to be the first thing in the file which conflicts with the RAT check. Make an exception for this file. (#1501) --- SKILL.md | 19 ------------------- dev/release/rat_exclude_files.txt | 3 ++- 2 files changed, 2 insertions(+), 20 deletions(-) diff --git a/SKILL.md b/SKILL.md index 9ba1c0cac..14ea5c609 100644 --- a/SKILL.md +++ b/SKILL.md @@ -1,22 +1,3 @@ - - --- name: datafusion-python description: Use when the user is writing datafusion-python (Apache DataFusion Python bindings) DataFrame or SQL code. Covers imports, data loading, DataFrame operations, expression building, SQL-to-DataFrame mappings, idiomatic patterns, and common pitfalls. diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index b2db144e8..a7a497dab 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -48,4 +48,5 @@ benchmarks/tpch/create_tables.sql .cargo/config.toml **/.cargo/config.toml uv.lock -examples/tpch/answers_sf1/*.tbl \ No newline at end of file +examples/tpch/answers_sf1/*.tbl +SKILL.md \ No newline at end of file From 8741d30cd812e4668f3f9187b56f12ce2de0d6e7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 23 Apr 2026 22:01:01 -0400 Subject: [PATCH 26/83] docs: enrich module docstrings and add doctest examples (#1498) * Enrich module docstrings and add doctest examples Expands the module docstrings for `functions.py`, `dataframe.py`, `expr.py`, and `context.py` so each module opens with a concept summary, cross-references to related APIs, and a small executable example. Adds doctest examples to the high-traffic `DataFrame` methods that previously lacked them: `select`, `aggregate`, `sort`, `limit`, `join`, and `union`. Optional parameters are demonstrated with keyword syntax, and examples reuse the same input data across variants so the effect of each option is easy to see. Co-Authored-By: Claude Opus 4.7 (1M context) * Use distinct group sums in aggregate docstring example Change the score data from [1, 2, 3] to [1, 2, 5] so the grouped result produces [3, 5] instead of [3, 3], removing ambiguity about which total belongs to which team. Co-Authored-By: Claude Opus 4.7 (1M context) * Align module-docstring examples with SKILL.md idioms Drop the redundant lit() in the dataframe.py module-docstring filter example and use a plain string group key in the aggregate() doctest, so both examples model the style SKILL.md recommends. Also document the sort("a") string form and sort_by() shortcut in SKILL.md's sorting section. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- SKILL.md | 16 +++- python/datafusion/context.py | 27 ++++++- python/datafusion/dataframe.py | 135 ++++++++++++++++++++++++++++++--- python/datafusion/expr.py | 28 ++++++- python/datafusion/functions.py | 22 +++++- 5 files changed, 209 insertions(+), 19 deletions(-) diff --git a/SKILL.md b/SKILL.md index 14ea5c609..7b07b430f 100644 --- a/SKILL.md +++ b/SKILL.md @@ -128,14 +128,22 @@ aggregate. ### Sorting ```python -df.sort(col("a")) # ascending (default) +df.sort("a") # ascending (plain name, preferred) +df.sort(col("a")) # ascending via col() df.sort(col("a").sort(ascending=False)) # descending df.sort(col("a").sort(nulls_first=False)) # override null placement + +df.sort_by("a", "b") # ascending-only shortcut ``` -A plain expression passed to `sort()` is already treated as ascending. Only -reach for `col(...).sort(...)` when you need to override a default (descending -order or null placement). Writing `col("a").sort(ascending=True)` is redundant. +As with `select()` and `aggregate()`, bare column references can be passed as +plain name strings. A plain expression passed to `sort()` is already treated +as ascending, so reach for `col(...).sort(...)` only when you need to override +a default (descending order or null placement). Writing +`col("a").sort(ascending=True)` is redundant. + +For ascending-only sorts with no null-placement override, `df.sort_by(...)` is +a shorter alias for `df.sort(...)`. ### Joining diff --git a/python/datafusion/context.py b/python/datafusion/context.py index c3f94cc16..dd6790402 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -15,7 +15,32 @@ # specific language governing permissions and limitations # under the License. -"""Session Context and it's associated configuration.""" +""":py:class:`SessionContext` — entry point for running DataFusion queries. + +A :py:class:`SessionContext` holds registered tables, catalogs, and +configuration for the current session. It is the first object most programs +create: from it you register data, run SQL strings +(:py:meth:`SessionContext.sql`), read files +(:py:meth:`SessionContext.read_csv`, +:py:meth:`SessionContext.read_parquet`, ...), and construct +:py:class:`~datafusion.dataframe.DataFrame` objects in memory +(:py:meth:`SessionContext.from_pydict`, +:py:meth:`SessionContext.from_arrow`). + +Session behavior (memory limits, batch size, configured optimizer passes, +...) is controlled by :py:class:`SessionConfig` and +:py:class:`RuntimeEnvBuilder`; SQL dialect limits are controlled by +:py:class:`SQLOptions`. + +Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> ctx.sql("SELECT 1 AS n").to_pydict() + {'n': [1]} + +See :ref:`user_guide_concepts` in the online documentation for the broader +execution model. +""" from __future__ import annotations diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index c00c85fdb..2b07861da 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -14,9 +14,32 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -""":py:class:`DataFrame` is one of the core concepts in DataFusion. - -See :ref:`user_guide_concepts` in the online documentation for more information. +""":py:class:`DataFrame` — lazy, chainable query representation. + +A :py:class:`DataFrame` is a logical plan over one or more data sources. +Methods that reshape the plan (:py:meth:`DataFrame.select`, +:py:meth:`DataFrame.filter`, :py:meth:`DataFrame.aggregate`, +:py:meth:`DataFrame.sort`, :py:meth:`DataFrame.join`, +:py:meth:`DataFrame.limit`, the set-operation methods, ...) return a new +:py:class:`DataFrame` and do no work until a terminal method such as +:py:meth:`DataFrame.collect`, :py:meth:`DataFrame.to_pydict`, +:py:meth:`DataFrame.show`, or one of the ``write_*`` methods is called. + +DataFrames are produced from a +:py:class:`~datafusion.context.SessionContext`, typically via +:py:meth:`~datafusion.context.SessionContext.sql`, +:py:meth:`~datafusion.context.SessionContext.read_csv`, +:py:meth:`~datafusion.context.SessionContext.read_parquet`, or +:py:meth:`~datafusion.context.SessionContext.from_pydict`. + +Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3], "b": [10, 20, 30]}) + >>> df.filter(col("a") > 1).select("b").to_pydict() + {'b': [20, 30]} + +See :ref:`user_guide_concepts` in the online documentation for a high-level +overview of the execution model. """ from __future__ import annotations @@ -503,21 +526,29 @@ def select_exprs(self, *args: str) -> DataFrame: def select(self, *exprs: Expr | str) -> DataFrame: """Project arbitrary expressions into a new :py:class:`DataFrame`. + String arguments are treated as column names; :py:class:`~datafusion.expr.Expr` + arguments can reshape, rename, or compute new columns. + Args: exprs: Either column names or :py:class:`~datafusion.expr.Expr` to select. Returns: DataFrame after projection. It has one column for each expression. - Example usage: + Examples: + Select columns by name: - The following example will return 3 columns from the original dataframe. - The first two columns will be the original column ``a`` and ``b`` since the - string "a" is assumed to refer to column selection. Also a duplicate of - column ``a`` will be returned with the column name ``alternate_a``:: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3], "b": [10, 20, 30]}) + >>> df.select("a").to_pydict() + {'a': [1, 2, 3]} - df = df.select("a", col("b"), col("a").alias("alternate_a")) + Mix column names, expressions, and aliases. The string ``"a"`` selects + column ``a`` directly; ``col("a").alias("alternate_a")`` returns a + duplicate under a new name: + >>> df.select("a", col("b"), col("a").alias("alternate_a")).to_pydict() + {'a': [1, 2, 3], 'b': [10, 20, 30], 'alternate_a': [1, 2, 3]} """ exprs_internal = expr_list_to_raw_expr_list(exprs) return DataFrame(self.df.select(*exprs_internal)) @@ -766,6 +797,24 @@ def aggregate( Returns: DataFrame after aggregation. + + Examples: + Aggregate without grouping — an empty ``group_by`` produces a + single row: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict( + ... {"team": ["x", "x", "y"], "score": [1, 2, 5]} + ... ) + >>> df.aggregate([], [F.sum(col("score")).alias("total")]).to_pydict() + {'total': [8]} + + Group by a column and produce one row per group: + + >>> df.aggregate( + ... ["team"], [F.sum(col("score")).alias("total")] + ... ).sort("team").to_pydict() + {'team': ['x', 'y'], 'total': [3, 5]} """ group_by_list = ( list(group_by) @@ -786,13 +835,27 @@ def sort(self, *exprs: SortKey) -> DataFrame: """Sort the DataFrame by the specified sorting expressions or column names. Note that any expression can be turned into a sort expression by - calling its ``sort`` method. + calling its ``sort`` method. For ascending-only sorts, the shorter + :py:meth:`sort_by` is usually more convenient. Args: exprs: Sort expressions or column names, applied in order. Returns: DataFrame after sorting. + + Examples: + Sort ascending by a column name: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [3, 1, 2], "b": [10, 20, 30]}) + >>> df.sort("a").to_pydict() + {'a': [1, 2, 3], 'b': [20, 30, 10]} + + Sort descending using :py:meth:`Expr.sort`: + + >>> df.sort(col("a").sort(ascending=False)).to_pydict() + {'a': [3, 2, 1], 'b': [10, 30, 20]} """ exprs_raw = sort_list_to_raw_sort_list(exprs) return DataFrame(self.df.sort(*exprs_raw)) @@ -812,12 +875,28 @@ def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame: def limit(self, count: int, offset: int = 0) -> DataFrame: """Return a new :py:class:`DataFrame` with a limited number of rows. + Results are returned in unspecified order unless the DataFrame is + explicitly sorted first via :py:meth:`sort` or :py:meth:`sort_by`. + Args: count: Number of rows to limit the DataFrame to. offset: Number of rows to skip. Returns: DataFrame after limiting. + + Examples: + Take the first two rows: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3, 4]}).sort("a") + >>> df.limit(2).to_pydict() + {'a': [1, 2]} + + Skip the first row then take two (paging): + + >>> df.limit(2, offset=1).to_pydict() + {'a': [2, 3]} """ return DataFrame(self.df.limit(count, offset)) @@ -972,6 +1051,28 @@ def join( Returns: DataFrame after join. + + Examples: + Inner-join two DataFrames on a shared column: + + >>> ctx = dfn.SessionContext() + >>> left = ctx.from_pydict({"id": [1, 2, 3], "val": [10, 20, 30]}) + >>> right = ctx.from_pydict({"id": [2, 3, 4], "label": ["b", "c", "d"]}) + >>> left.join(right, on="id").sort("id").to_pydict() + {'id': [2, 3], 'val': [20, 30], 'label': ['b', 'c']} + + Left join to keep all rows from the left side: + + >>> left.join(right, on="id", how="left").sort("id").to_pydict() + {'id': [1, 2, 3], 'val': [10, 20, 30], 'label': [None, 'b', 'c']} + + Use ``left_on`` / ``right_on`` when the key columns differ in name: + + >>> right2 = ctx.from_pydict({"rid": [2, 3], "label": ["b", "c"]}) + >>> left.join( + ... right2, left_on="id", right_on="rid" + ... ).sort("id").to_pydict() + {'id': [2, 3], 'val': [20, 30], 'rid': [2, 3], 'label': ['b', 'c']} """ if join_keys is not None: warnings.warn( @@ -1165,6 +1266,20 @@ def union(self, other: DataFrame, distinct: bool = False) -> DataFrame: Returns: DataFrame after union. + + Examples: + Stack rows from both DataFrames, preserving duplicates: + + >>> ctx = dfn.SessionContext() + >>> df1 = ctx.from_pydict({"a": [1, 2]}) + >>> df2 = ctx.from_pydict({"a": [2, 3]}) + >>> df1.union(df2).sort("a").to_pydict() + {'a': [1, 2, 2, 3]} + + Deduplicate the combined result with ``distinct=True``: + + >>> df1.union(df2, distinct=True).sort("a").to_pydict() + {'a': [1, 2, 3]} """ return DataFrame(self.df.union(other.df, distinct)) diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 32004656f..1ff6976f7 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -15,9 +15,31 @@ # specific language governing permissions and limitations # under the License. -"""This module supports expressions, one of the core concepts in DataFusion. - -See :ref:`Expressions` in the online documentation for more details. +""":py:class:`Expr` — the logical expression type used to build DataFusion queries. + +An :py:class:`Expr` represents a computation over columns or literals: a +column reference (``col("a")``), a literal (``lit(5)``), an operator +combination (``col("a") + lit(1)``), or the output of a function from +:py:mod:`datafusion.functions`. Expressions are passed to +:py:class:`~datafusion.dataframe.DataFrame` methods such as +:py:meth:`~datafusion.dataframe.DataFrame.select`, +:py:meth:`~datafusion.dataframe.DataFrame.filter`, +:py:meth:`~datafusion.dataframe.DataFrame.aggregate`, and +:py:meth:`~datafusion.dataframe.DataFrame.sort`. + +Convenience constructors are re-exported at the package level: +:py:func:`datafusion.col` / :py:func:`datafusion.column` for column references +and :py:func:`datafusion.lit` / :py:func:`datafusion.literal` for scalar +literals. + +Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> df.select((col("a") * lit(10)).alias("ten_a")).to_pydict() + {'ten_a': [10, 20, 30]} + +See :ref:`expressions` in the online documentation for details on available +operators and helpers. """ # ruff: noqa: PLC0415 diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 841cd9c0b..280a6d3ac 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -14,7 +14,27 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""User functions for operating on :py:class:`~datafusion.expr.Expr`.""" +"""Scalar, aggregate, and window functions for :py:class:`~datafusion.expr.Expr`. + +Each function returns an :py:class:`~datafusion.expr.Expr` that can be combined +with other expressions and passed to +:py:class:`~datafusion.dataframe.DataFrame` methods such as +:py:meth:`~datafusion.dataframe.DataFrame.select`, +:py:meth:`~datafusion.dataframe.DataFrame.filter`, +:py:meth:`~datafusion.dataframe.DataFrame.aggregate`, and +:py:meth:`~datafusion.dataframe.DataFrame.window`. The module is conventionally +imported as ``F`` so calls read like ``F.sum(col("price"))``. + +Examples: + >>> from datafusion import functions as F + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3, 4]}) + >>> df.aggregate([], [F.sum(col("a")).alias("total")]).to_pydict() + {'total': [10]} + +See :ref:`aggregation` and :ref:`window_functions` in the online documentation +for categorized catalogs of aggregate and window functions. +""" from __future__ import annotations From c8bb9f7d3876de97141d204740a6b99d5facd10f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 24 Apr 2026 07:57:11 -0400 Subject: [PATCH 27/83] docs: add README section for AI coding assistants (#1503) Points users to the repo-root SKILL.md via the npx skills registry or a manual AGENTS.md / CLAUDE.md pointer. Implements PR 1c of the plan in #1394. Co-authored-by: Claude Opus 4.7 (1M context) --- README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/README.md b/README.md index 7849e7a02..4baed7d1d 100644 --- a/README.md +++ b/README.md @@ -215,6 +215,22 @@ You can verify the installation by running: '0.6.0' ``` +## 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 +[Agent Skills](https://agentskills.io) open standard. + +**Preferred:** `npx skills add apache/datafusion-python` — installs the skill in +Claude Code, Cursor, Windsurf, Cline, Codex, Copilot, Gemini CLI, and other +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 +``` + ## How to develop This assumes that you have rust and cargo installed. We use the workflow recommended by [pyo3](https://github.com/PyO3/pyo3) and [maturin](https://github.com/PyO3/maturin). The Maturin tools used in this workflow can be installed either via `uv` or `pip`. Both approaches should offer the same experience. It is recommended to use `uv` since it has significant performance improvements From 03577163a057f791b19f30ce5130464a4a1c78a4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 24 Apr 2026 11:47:06 -0400 Subject: [PATCH 28/83] tpch examples: rewrite queries idiomatically and embed reference SQL (#1504) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * tpch examples: add reference SQL to each query, fix Q20 - Append the canonical TPC-H reference SQL (from benchmarks/tpch/queries/) to each q01..q22 module docstring so readers can compare the DataFrame translation against the SQL at a glance. - Fix Q20: `df = df.filter(col("ps_availqty") > lit(0.5) * col("total_sold"))` was missing the assignment so the filter was dropped from the pipeline. Co-Authored-By: Claude Opus 4.7 (1M context) * tpch examples: rewrite non-idiomatic queries in idiomatic DataFrame form Rewrite the seven TPC-H example queries that did not demonstrate the idiomatic DataFrame pattern. The remaining queries (Q02/Q11/Q15/Q17/Q22, which use window functions in place of correlated subqueries) already are idiomatic and are left unchanged. - Q04: replace `.aggregate([col("l_orderkey")], [])` with `.select("l_orderkey").distinct()`, which is the natural way to express "reduce to one row per order" on a DataFrame. - Q07: remove the CASE-as-filter on `n_name` and use `F.in_list(col("n_name"), [nation_1, nation_2])` instead. Drops a comment block that admitted the filter form was simpler. - Q08: rewrite the switched CASE `F.case(...).when(lit(False), ...)` as a searched `F.when(col(...).is_not_null(), ...).otherwise(...)`. That mirrors the reference SQL's `case when ... then ... else 0 end` shape. - Q12: replace `array_position(make_array(...), col)` with `F.in_list(col("l_shipmode"), [...])`. Same semantics, without routing through array construction / array search. - Q19: remove the pyarrow UDF that re-implemented a disjunctive predicate in Python. Build the same predicate in DataFusion by OR-combining one `in_list` + range-filter expression per brand. Keeps the per-brand constants in the existing `items_of_interest` dict. - Q20: use `F.starts_with` instead of an explicit substring slice. Replace the inner-join + `select(...).distinct()` tail with a semi join against a precomputed set of excess-quantity suppliers so the supplier columns are preserved without deduplication after the fact. - Q21: replace the `array_agg` / `array_length` / `array_element` pipeline with two semi joins. One semi join keeps orders with more than one distinct supplier (stand-in for the reference SQL's `exists` subquery), the other keeps orders with exactly one late supplier (stand-in for the `not exists` subquery). All 22 answer-file comparisons and 22 plan-comparison diagnostics still pass (`pytest examples/tpch/_tests.py`: 44 passed). Co-Authored-By: Claude Opus 4.7 (1M context) * tpch examples: align reference SQL constants with DataFrame queries The reference SQL embedded in each q01..q22 module docstring was carried over verbatim from ``benchmarks/tpch/queries/`` and uses a different set of TPC-H substitution parameters than the DataFrame examples (answer-file-validated at scale factor 1). Update each reference SQL to use the substitution parameters the DataFrame uses, so both expressions describe the same query and would produce the same results against the same data. Constants aligned: - Q01: ``90 days`` cutoff (DataFrame ``DAYS_BEFORE_FINAL = 90``). - Q02: ``p_size = 15``, ``p_type like '%BRASS'``, ``r_name = 'EUROPE'``. - Q04: base date ``1993-07-01`` (``3 month`` interval preserved per the "quarter of a year" wording). - Q05: ``r_name = 'ASIA'``. - Q06: ``l_discount between 0.06 - 0.01 and 0.06 + 0.01``. - Q07: nations ``'FRANCE'`` / ``'GERMANY'``. - Q08: ``r_name = 'AMERICA'``, ``p_type = 'ECONOMY ANODIZED STEEL'``, inner-case ``nation = 'BRAZIL'``. - Q09: ``p_name like '%green%'``. - Q10: base date ``1993-10-01`` (``3 month`` interval preserved). - Q11: ``n_name = 'GERMANY'``. - Q12: ship modes ``('MAIL', 'SHIP')``, base date ``1994-01-01``. - Q13: ``o_comment not like '%special%requests%'``. - Q14: base date ``1995-09-01``. - Q15: base date ``1996-01-01``. - Q16: ``p_brand <> 'Brand#45'``, ``p_type not like 'MEDIUM POLISHED%'``, sizes ``(49, 14, 23, 45, 19, 3, 36, 9)``. - Q17: ``p_brand = 'Brand#23'``, ``p_container = 'MED BOX'``. - Q18: ``sum(l_quantity) > 300``. - Q19: brands ``Brand#12`` / ``Brand#23`` / ``Brand#34`` with the matching minimum quantities (1, 10, 20). - Q20: ``p_name like 'forest%'``, base date ``1994-01-01``, ``n_name = 'CANADA'``. - Q21: ``n_name = 'SAUDI ARABIA'``. - Q22: country codes ``('13', '31', '23', '29', '30', '18', '17')``. Interval units (month / year) are preserved where the problem-statement text reads "given quarter", "given year", "given month". Q01 keeps the literal "days" unit because the TPC-H problem statement itself describes the cutoff in days. Co-Authored-By: Claude Opus 4.7 (1M context) * tpch examples: apply SKILL.md idioms across all 22 queries Sweep every q01..q22 example for idiomatic DataFrame style as described in the repo-root SKILL.md: - ``col("x") == "s"`` in place of ``col("x") == lit("s")`` on comparison right-hand sides (auto-wrap applies). - Plain-name strings in ``select``/``aggregate``/``sort`` group/sort key lists when the key is a bare column. - Drop redundant ``how="inner"`` and single-element ``left_on``/``right_on`` list wrapping on equi-joins. - Collapse chained ``.filter(a).filter(b)`` runs into ``.filter(a, b)`` and chained ``.with_column`` runs into ``.with_columns(a=..., b=...)``. - ``df.sort_by(...)`` or plain-name ``df.sort(...)`` when no null-placement override is needed. - ``F.count_star()`` in place of ``F.count(col("x"))`` whenever the SQL reads ``count(*)``. - ``F.starts_with(col, lit(prefix))`` and ``~F.starts_with(...)`` in place of substring-prefix equality/inequality tricks. - ``F.in_list(col, [lit(...)])`` in place of ``~F.array_position(...). is_null()`` and in place of disjunctions of equality comparisons. - Searched ``F.when(cond, x).otherwise(y)`` in place of switched ``F.case(bool_expr).when(lit(True/False), x).end()`` forms. - Semi-joins as the DataFrame form of ``EXISTS`` (Q04); anti-joins as ``NOT EXISTS`` (Q22 was already using this idiom). - Whole-frame window aggregates as the DataFrame stand-in for a SQL scalar subquery (Q11/Q15/Q17/Q22). Individual query fixes of note: - Q16 — add the secondary sort keys (``p_brand``, ``p_type``, ``p_size``) that the TPC-H spec requires but the original DataFrame omitted. - Q22 — drop a stray ``df.show()`` mid-pipeline; replace the 0-based substring slice with ``F.left(col("c_phone"), lit(2))``. - Q14 — rewrite the promo/non-promo factor split as a searched CASE inside ``F.sum(...)`` so the DataFrame expression matches the reference SQL shape exactly. All 22 answer-file comparisons still pass at scale factor 1. Co-Authored-By: Claude Opus 4.7 (1M context) * tpch examples: more idiomatic aggregate FILTER, string funcs, date handling Additional sweep of the TPC-H DataFrame examples informed by comparing against a fresh set of SKILL.md-only generations under ``examples/tpch/agentic_queries/``: - Q02: ``F.ends_with(col("p_type"), lit(TYPE_OF_INTEREST))`` in place of ``F.strpos(col, lit) > 0``. The reference SQL is ``p_type like '%BRASS'``, which is an ends_with check, not contains. ``F.strpos > 0`` returned the correct rows on TPC-H data by coincidence but is semantically wrong. - Q09: ``F.contains(col("p_name"), lit(part_color))`` in place of ``F.strpos(col, lit) > 0``. The SQL is ``p_name like '%green%'``. - Q08, Q12, Q14: use the ``filter`` keyword on ``F.sum`` / ``F.count`` — the DataFrame form of SQL ``sum(...) FILTER (WHERE ...)`` — instead of wrapping the aggregate input in ``F.when(cond, x).otherwise(0)``. Q08 also reorganises to inner-join the supplier's nation onto the regional sales, which removes the previous left-join + ``F.when(is_not_null, ...)`` dance. - Q15: compute the grand maximum revenue as a separate scalar aggregate and ``join_on(...)`` on equality, instead of the whole-frame window ``F.max`` + filter shape. Simpler plan, same result. - Q16: ``F.regexp_like(col, pattern)`` in place of ``F.regexp_match(col, pattern).is_not_null()``. - Q04, Q05, Q06, Q07, Q08, Q10, Q12, Q14, Q15, Q20: store both the start and the end of the date window as plain ``datetime.date`` objects and compare with ``lit(end_date)``, instead of carrying the start date + ``pa.month_day_nano_interval`` and adding them at query-build time. Drops unused ``pyarrow`` imports from the files that no longer need Arrow scalars. All 22 answer-file comparisons still pass at scale factor 1. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- examples/tpch/q01_pricing_summary_report.py | 44 ++-- examples/tpch/q02_minimum_cost_supplier.py | 87 ++++++-- examples/tpch/q03_shipping_priority.py | 51 +++-- examples/tpch/q04_order_priority_checking.py | 67 +++--- examples/tpch/q05_local_supplier_volume.py | 66 +++--- .../tpch/q06_forecasting_revenue_change.py | 35 +-- examples/tpch/q07_volume_shipping.py | 103 +++++---- examples/tpch/q08_market_share.py | 205 +++++++++--------- .../tpch/q09_product_type_profit_measure.py | 77 +++++-- examples/tpch/q10_returned_item_reporting.py | 102 +++++---- .../q11_important_stock_identification.py | 83 ++++--- examples/tpch/q12_ship_mode_order_priority.py | 108 +++++---- examples/tpch/q13_customer_distribution.py | 47 ++-- examples/tpch/q14_promotion_effect.py | 81 +++---- examples/tpch/q15_top_supplier.py | 94 ++++---- .../tpch/q16_part_supplier_relationship.py | 81 ++++--- examples/tpch/q17_small_quantity_order.py | 58 +++-- examples/tpch/q18_large_volume_customer.py | 71 ++++-- examples/tpch/q19_discounted_revenue.py | 134 ++++++------ examples/tpch/q20_potential_part_promotion.py | 120 ++++++---- .../tpch/q21_suppliers_kept_orders_waiting.py | 134 +++++++----- examples/tpch/q22_global_sales_opportunity.py | 104 ++++++--- 22 files changed, 1196 insertions(+), 756 deletions(-) diff --git a/examples/tpch/q01_pricing_summary_report.py b/examples/tpch/q01_pricing_summary_report.py index 3f97f00dc..105f1632d 100644 --- a/examples/tpch/q01_pricing_summary_report.py +++ b/examples/tpch/q01_pricing_summary_report.py @@ -27,6 +27,30 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + l_returnflag, + l_linestatus, + sum(l_quantity) as sum_qty, + sum(l_extendedprice) as sum_base_price, + sum(l_extendedprice * (1 - l_discount)) as sum_disc_price, + sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge, + avg(l_quantity) as avg_qty, + avg(l_extendedprice) as avg_price, + avg(l_discount) as avg_disc, + count(*) as count_order + from + lineitem + where + l_shipdate <= date '1998-12-01' - interval '90 days' + group by + l_returnflag, + l_linestatus + order by + l_returnflag, + l_linestatus; """ import pyarrow as pa @@ -58,31 +82,25 @@ # Aggregate the results +disc_price = col("l_extendedprice") * (lit(1) - col("l_discount")) + df = df.aggregate( - [col("l_returnflag"), col("l_linestatus")], + ["l_returnflag", "l_linestatus"], [ F.sum(col("l_quantity")).alias("sum_qty"), F.sum(col("l_extendedprice")).alias("sum_base_price"), - F.sum(col("l_extendedprice") * (lit(1) - col("l_discount"))).alias( - "sum_disc_price" - ), - F.sum( - col("l_extendedprice") - * (lit(1) - col("l_discount")) - * (lit(1) + col("l_tax")) - ).alias("sum_charge"), + F.sum(disc_price).alias("sum_disc_price"), + F.sum(disc_price * (lit(1) + col("l_tax"))).alias("sum_charge"), F.avg(col("l_quantity")).alias("avg_qty"), F.avg(col("l_extendedprice")).alias("avg_price"), F.avg(col("l_discount")).alias("avg_disc"), - F.count(col("l_returnflag")).alias( - "count_order" - ), # Counting any column should return same result + F.count_star().alias("count_order"), ], ) # Sort per the expected result -df = df.sort(col("l_returnflag").sort(), col("l_linestatus").sort()) +df = df.sort_by("l_returnflag", "l_linestatus") # Note: There appears to be a discrepancy between what is returned here and what is in the generated # answers file for the case of return flag N and line status O, but I did not investigate further. diff --git a/examples/tpch/q02_minimum_cost_supplier.py b/examples/tpch/q02_minimum_cost_supplier.py index 47961d2ef..c5c6b9c0b 100644 --- a/examples/tpch/q02_minimum_cost_supplier.py +++ b/examples/tpch/q02_minimum_cost_supplier.py @@ -27,6 +27,52 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + s_acctbal, + s_name, + n_name, + p_partkey, + p_mfgr, + s_address, + s_phone, + s_comment + from + part, + supplier, + partsupp, + nation, + region + where + p_partkey = ps_partkey + and s_suppkey = ps_suppkey + and p_size = 15 + and p_type like '%BRASS' + and s_nationkey = n_nationkey + and n_regionkey = r_regionkey + and r_name = 'EUROPE' + and ps_supplycost = ( + select + min(ps_supplycost) + from + partsupp, + supplier, + nation, + region + where + p_partkey = ps_partkey + and s_suppkey = ps_suppkey + and s_nationkey = n_nationkey + and n_regionkey = r_regionkey + and r_name = 'EUROPE' + ) + order by + s_acctbal desc, + n_name, + s_name, + p_partkey limit 100; """ import datafusion @@ -67,35 +113,30 @@ "r_regionkey", "r_name" ) -# Filter down parts. Part names contain the type of interest, so we can use strpos to find where -# in the p_type column the word is. `strpos` will return 0 if not found, otherwise the position -# in the string where it is located. +# Filter down parts. The reference SQL uses ``p_type like '%BRASS'`` which +# is an ``ends_with`` check; use the dedicated string function rather than +# a manual substring match. df_part = df_part.filter( - F.strpos(col("p_type"), lit(TYPE_OF_INTEREST)) > lit(0) -).filter(col("p_size") == lit(SIZE_OF_INTEREST)) + F.ends_with(col("p_type"), lit(TYPE_OF_INTEREST)), + col("p_size") == SIZE_OF_INTEREST, +) # Filter regions down to the one of interest -df_region = df_region.filter(col("r_name") == lit(REGION_OF_INTEREST)) +df_region = df_region.filter(col("r_name") == REGION_OF_INTEREST) # Now that we have the region, find suppliers in that region. Suppliers are tied to their nation # and nations are tied to the region. -df_nation = df_nation.join( - df_region, left_on=["n_regionkey"], right_on=["r_regionkey"], how="inner" -) -df_supplier = df_supplier.join( - df_nation, left_on=["s_nationkey"], right_on=["n_nationkey"], how="inner" -) +df_nation = df_nation.join(df_region, left_on="n_regionkey", right_on="r_regionkey") +df_supplier = df_supplier.join(df_nation, left_on="s_nationkey", right_on="n_nationkey") # Now that we know who the potential suppliers are for the part, we can limit out part # supplies table down. We can further join down to the specific parts we've identified # as matching the request -df = df_partsupp.join( - df_supplier, left_on=["ps_suppkey"], right_on=["s_suppkey"], how="inner" -) +df = df_partsupp.join(df_supplier, left_on="ps_suppkey", right_on="s_suppkey") # Locate the minimum cost across all suppliers. There are multiple ways you could do this, # but one way is to create a window function across all suppliers, find the minimum, and @@ -112,9 +153,9 @@ ), ) -df = df.filter(col("min_cost") == col("ps_supplycost")) - -df = df.join(df_part, left_on=["ps_partkey"], right_on=["p_partkey"], how="inner") +df = df.filter(col("min_cost") == col("ps_supplycost")).join( + df_part, left_on="ps_partkey", right_on="p_partkey" +) # From the problem statement, these are the values we wish to output @@ -132,12 +173,10 @@ # Sort and display 100 entries df = df.sort( col("s_acctbal").sort(ascending=False), - col("n_name").sort(), - col("s_name").sort(), - col("p_partkey").sort(), -) - -df = df.limit(100) + "n_name", + "s_name", + "p_partkey", +).limit(100) # Show results diff --git a/examples/tpch/q03_shipping_priority.py b/examples/tpch/q03_shipping_priority.py index fc1231e0a..880c7435f 100644 --- a/examples/tpch/q03_shipping_priority.py +++ b/examples/tpch/q03_shipping_priority.py @@ -25,6 +25,31 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + l_orderkey, + sum(l_extendedprice * (1 - l_discount)) as revenue, + o_orderdate, + o_shippriority + from + customer, + orders, + lineitem + where + c_mktsegment = 'BUILDING' + and c_custkey = o_custkey + and l_orderkey = o_orderkey + and o_orderdate < date '1995-03-15' + and l_shipdate > date '1995-03-15' + group by + l_orderkey, + o_orderdate, + o_shippriority + order by + revenue desc, + o_orderdate limit 10; """ from datafusion import SessionContext, col, lit @@ -50,20 +75,20 @@ # Limit dataframes to the rows of interest -df_customer = df_customer.filter(col("c_mktsegment") == lit(SEGMENT_OF_INTEREST)) +df_customer = df_customer.filter(col("c_mktsegment") == SEGMENT_OF_INTEREST) df_orders = df_orders.filter(col("o_orderdate") < lit(DATE_OF_INTEREST)) df_lineitem = df_lineitem.filter(col("l_shipdate") > lit(DATE_OF_INTEREST)) # Join all 3 dataframes -df = df_customer.join( - df_orders, left_on=["c_custkey"], right_on=["o_custkey"], how="inner" -).join(df_lineitem, left_on=["o_orderkey"], right_on=["l_orderkey"], how="inner") +df = df_customer.join(df_orders, left_on="c_custkey", right_on="o_custkey").join( + df_lineitem, left_on="o_orderkey", right_on="l_orderkey" +) # Compute the revenue df = df.aggregate( - [col("l_orderkey")], + ["l_orderkey"], [ F.first_value(col("o_orderdate")).alias("o_orderdate"), F.first_value(col("o_shippriority")).alias("o_shippriority"), @@ -71,17 +96,13 @@ ], ) -# Sort by priority - -df = df.sort(col("revenue").sort(ascending=False), col("o_orderdate").sort()) - -# Only return 10 results +# Sort by priority, take 10, and project in the order expected by the spec. -df = df.limit(10) - -# Change the order that the columns are reported in just to match the spec - -df = df.select("l_orderkey", "revenue", "o_orderdate", "o_shippriority") +df = ( + df.sort(col("revenue").sort(ascending=False), "o_orderdate") + .limit(10) + .select("l_orderkey", "revenue", "o_orderdate", "o_shippriority") +) # Show result diff --git a/examples/tpch/q04_order_priority_checking.py b/examples/tpch/q04_order_priority_checking.py index 426338aea..6f11c1383 100644 --- a/examples/tpch/q04_order_priority_checking.py +++ b/examples/tpch/q04_order_priority_checking.py @@ -24,18 +24,40 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + o_orderpriority, + count(*) as order_count + from + orders + where + o_orderdate >= date '1993-07-01' + and o_orderdate < date '1993-07-01' + interval '3' month + and exists ( + select + * + from + lineitem + where + l_orderkey = o_orderkey + and l_commitdate < l_receiptdate + ) + group by + o_orderpriority + order by + o_orderpriority; """ -from datetime import datetime +from datetime import date -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path -# Ideally we could put 3 months into the interval. See note below. -INTERVAL_DAYS = 92 -DATE_OF_INTEREST = "1993-07-01" +QUARTER_START = date(1993, 7, 1) +QUARTER_END = date(1993, 10, 1) # Load the dataframes we need @@ -48,36 +70,23 @@ "l_orderkey", "l_commitdate", "l_receiptdate" ) -# Create a date object from the string -date = datetime.strptime(DATE_OF_INTEREST, "%Y-%m-%d").date() - -interval = pa.scalar((0, INTERVAL_DAYS, 0), type=pa.month_day_nano_interval()) - -# Limit results to cases where commitment date before receipt date -# Aggregate the results so we only get one row to join with the order table. -# Alternately, and likely more idiomatic is instead of `.aggregate` you could -# do `.select("l_orderkey").distinct()`. The goal here is to show -# multiple examples of how to use Data Fusion. -df_lineitem = df_lineitem.filter(col("l_commitdate") < col("l_receiptdate")).aggregate( - [col("l_orderkey")], [] +# Keep only orders in the quarter of interest, then restrict to those that +# have at least one late lineitem via a semi join (the DataFrame form of +# ``EXISTS`` from the reference SQL). +df_orders = df_orders.filter( + col("o_orderdate") >= lit(QUARTER_START), + col("o_orderdate") < lit(QUARTER_END), ) -# Limit orders to date range of interest -df_orders = df_orders.filter(col("o_orderdate") >= lit(date)).filter( - col("o_orderdate") < lit(date) + lit(interval) -) +late_lineitems = df_lineitem.filter(col("l_commitdate") < col("l_receiptdate")) -# Perform the join to find only orders for which there are lineitems outside of expected range df = df_orders.join( - df_lineitem, left_on=["o_orderkey"], right_on=["l_orderkey"], how="inner" + late_lineitems, left_on="o_orderkey", right_on="l_orderkey", how="semi" ) -# Based on priority, find the number of entries -df = df.aggregate( - [col("o_orderpriority")], [F.count(col("o_orderpriority")).alias("order_count")] +# Count the number of orders in each priority group and sort. +df = df.aggregate(["o_orderpriority"], [F.count_star().alias("order_count")]).sort_by( + "o_orderpriority" ) -# Sort the results -df = df.sort(col("o_orderpriority").sort()) - df.show() diff --git a/examples/tpch/q05_local_supplier_volume.py b/examples/tpch/q05_local_supplier_volume.py index fa2b01dea..bfdba5d4c 100644 --- a/examples/tpch/q05_local_supplier_volume.py +++ b/examples/tpch/q05_local_supplier_volume.py @@ -27,23 +27,45 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + n_name, + sum(l_extendedprice * (1 - l_discount)) as revenue + from + customer, + orders, + lineitem, + supplier, + nation, + region + where + c_custkey = o_custkey + and l_orderkey = o_orderkey + and l_suppkey = s_suppkey + and c_nationkey = s_nationkey + and s_nationkey = n_nationkey + and n_regionkey = r_regionkey + and r_name = 'ASIA' + and o_orderdate >= date '1994-01-01' + and o_orderdate < date '1994-01-01' + interval '1' year + group by + n_name + order by + revenue desc; """ -from datetime import datetime +from datetime import date -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path -DATE_OF_INTEREST = "1994-01-01" -INTERVAL_DAYS = 365 +YEAR_START = date(1994, 1, 1) +YEAR_END = date(1995, 1, 1) REGION_OF_INTEREST = "ASIA" -date = datetime.strptime(DATE_OF_INTEREST, "%Y-%m-%d").date() - -interval = pa.scalar((0, INTERVAL_DAYS, 0), type=pa.month_day_nano_interval()) - # Load the dataframes we need ctx = SessionContext() @@ -68,38 +90,32 @@ ) # Restrict dataframes to cases of interest -df_orders = df_orders.filter(col("o_orderdate") >= lit(date)).filter( - col("o_orderdate") < lit(date) + lit(interval) +df_orders = df_orders.filter( + col("o_orderdate") >= lit(YEAR_START), + col("o_orderdate") < lit(YEAR_END), ) -df_region = df_region.filter(col("r_name") == lit(REGION_OF_INTEREST)) +df_region = df_region.filter(col("r_name") == REGION_OF_INTEREST) # Join all the dataframes df = ( - df_customer.join( - df_orders, left_on=["c_custkey"], right_on=["o_custkey"], how="inner" - ) - .join(df_lineitem, left_on=["o_orderkey"], right_on=["l_orderkey"], how="inner") + df_customer.join(df_orders, left_on="c_custkey", right_on="o_custkey") + .join(df_lineitem, left_on="o_orderkey", right_on="l_orderkey") .join( df_supplier, left_on=["l_suppkey", "c_nationkey"], right_on=["s_suppkey", "s_nationkey"], - how="inner", ) - .join(df_nation, left_on=["s_nationkey"], right_on=["n_nationkey"], how="inner") - .join(df_region, left_on=["n_regionkey"], right_on=["r_regionkey"], how="inner") + .join(df_nation, left_on="s_nationkey", right_on="n_nationkey") + .join(df_region, left_on="n_regionkey", right_on="r_regionkey") ) -# Compute the final result +# Compute the final result, then sort in descending order. df = df.aggregate( - [col("n_name")], + ["n_name"], [F.sum(col("l_extendedprice") * (lit(1.0) - col("l_discount"))).alias("revenue")], -) - -# Sort in descending order - -df = df.sort(col("revenue").sort(ascending=False)) +).sort(col("revenue").sort(ascending=False)) df.show() diff --git a/examples/tpch/q06_forecasting_revenue_change.py b/examples/tpch/q06_forecasting_revenue_change.py index 1de5848b1..ed54d22a4 100644 --- a/examples/tpch/q06_forecasting_revenue_change.py +++ b/examples/tpch/q06_forecasting_revenue_change.py @@ -27,28 +27,34 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + sum(l_extendedprice * l_discount) as revenue + from + lineitem + where + l_shipdate >= date '1994-01-01' + and l_shipdate < date '1994-01-01' + interval '1' year + and l_discount between 0.06 - 0.01 and 0.06 + 0.01 + and l_quantity < 24; """ -from datetime import datetime +from datetime import date -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path # Variables from the example query -DATE_OF_INTEREST = "1994-01-01" +YEAR_START = date(1994, 1, 1) +YEAR_END = date(1995, 1, 1) DISCOUT = 0.06 DELTA = 0.01 QUANTITY = 24 -INTERVAL_DAYS = 365 - -date = datetime.strptime(DATE_OF_INTEREST, "%Y-%m-%d").date() - -interval = pa.scalar((0, INTERVAL_DAYS, 0), type=pa.month_day_nano_interval()) - # Load the dataframes we need ctx = SessionContext() @@ -59,12 +65,11 @@ # Filter down to lineitems of interest -df = ( - df_lineitem.filter(col("l_shipdate") >= lit(date)) - .filter(col("l_shipdate") < lit(date) + lit(interval)) - .filter(col("l_discount") >= lit(DISCOUT) - lit(DELTA)) - .filter(col("l_discount") <= lit(DISCOUT) + lit(DELTA)) - .filter(col("l_quantity") < lit(QUANTITY)) +df = df_lineitem.filter( + col("l_shipdate") >= lit(YEAR_START), + col("l_shipdate") < lit(YEAR_END), + col("l_discount").between(lit(DISCOUT - DELTA), lit(DISCOUT + DELTA)), + col("l_quantity") < QUANTITY, ) # Add up all the "lost" revenue diff --git a/examples/tpch/q07_volume_shipping.py b/examples/tpch/q07_volume_shipping.py index ff2f891f1..df1c2ae0d 100644 --- a/examples/tpch/q07_volume_shipping.py +++ b/examples/tpch/q07_volume_shipping.py @@ -26,9 +26,51 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + supp_nation, + cust_nation, + l_year, + sum(volume) as revenue + from + ( + select + n1.n_name as supp_nation, + n2.n_name as cust_nation, + extract(year from l_shipdate) as l_year, + l_extendedprice * (1 - l_discount) as volume + from + supplier, + lineitem, + orders, + customer, + nation n1, + nation n2 + where + s_suppkey = l_suppkey + and o_orderkey = l_orderkey + and c_custkey = o_custkey + and s_nationkey = n1.n_nationkey + and c_nationkey = n2.n_nationkey + and ( + (n1.n_name = 'FRANCE' and n2.n_name = 'GERMANY') + or (n1.n_name = 'GERMANY' and n2.n_name = 'FRANCE') + ) + and l_shipdate between date '1995-01-01' and date '1996-12-31' + ) as shipping + group by + supp_nation, + cust_nation, + l_year + order by + supp_nation, + cust_nation, + l_year; """ -from datetime import datetime +from datetime import date import pyarrow as pa from datafusion import SessionContext, col, lit @@ -40,11 +82,8 @@ nation_1 = lit("FRANCE") nation_2 = lit("GERMANY") -START_DATE = "1995-01-01" -END_DATE = "1996-12-31" - -start_date = lit(datetime.strptime(START_DATE, "%Y-%m-%d").date()) -end_date = lit(datetime.strptime(END_DATE, "%Y-%m-%d").date()) +START_DATE = date(1995, 1, 1) +END_DATE = date(1996, 12, 31) # Load the dataframes we need @@ -69,60 +108,44 @@ # Filter to time of interest -df_lineitem = df_lineitem.filter(col("l_shipdate") >= start_date).filter( - col("l_shipdate") <= end_date +df_lineitem = df_lineitem.filter( + col("l_shipdate") >= lit(START_DATE), col("l_shipdate") <= lit(END_DATE) ) -# A simpler way to do the following operation is to use a filter, but we also want to demonstrate -# how to use case statements. Here we are assigning `n_name` to be itself when it is either of -# the two nations of interest. Since there is no `otherwise()` statement, any values that do -# not match these will result in a null value and then get filtered out. -# -# To do the same using a simple filter would be: -# df_nation = df_nation.filter((F.col("n_name") == nation_1) | (F.col("n_name") == nation_2)) # noqa: ERA001 -df_nation = df_nation.with_column( - "n_name", - F.case(col("n_name")) - .when(nation_1, col("n_name")) - .when(nation_2, col("n_name")) - .end(), -).filter(~col("n_name").is_null()) +# Limit the nation table to the two nations of interest. +df_nation = df_nation.filter(F.in_list(col("n_name"), [nation_1, nation_2])) # Limit suppliers to either nation df_supplier = df_supplier.join( - df_nation, left_on=["s_nationkey"], right_on=["n_nationkey"], how="inner" -).select(col("s_suppkey"), col("n_name").alias("supp_nation")) + df_nation, left_on="s_nationkey", right_on="n_nationkey" +).select("s_suppkey", col("n_name").alias("supp_nation")) # Limit customers to either nation df_customer = df_customer.join( - df_nation, left_on=["c_nationkey"], right_on=["n_nationkey"], how="inner" -).select(col("c_custkey"), col("n_name").alias("cust_nation")) + df_nation, left_on="c_nationkey", right_on="n_nationkey" +).select("c_custkey", col("n_name").alias("cust_nation")) # Join up all the data frames from line items, and make sure the supplier and customer are in # different nations. df = ( - df_lineitem.join( - df_orders, left_on=["l_orderkey"], right_on=["o_orderkey"], how="inner" - ) - .join(df_customer, left_on=["o_custkey"], right_on=["c_custkey"], how="inner") - .join(df_supplier, left_on=["l_suppkey"], right_on=["s_suppkey"], how="inner") + df_lineitem.join(df_orders, left_on="l_orderkey", right_on="o_orderkey") + .join(df_customer, left_on="o_custkey", right_on="c_custkey") + .join(df_supplier, left_on="l_suppkey", right_on="s_suppkey") .filter(col("cust_nation") != col("supp_nation")) ) # Extract out two values for every line item -df = df.with_column( - "l_year", F.datepart(lit("year"), col("l_shipdate")).cast(pa.int32()) -).with_column("volume", col("l_extendedprice") * (lit(1.0) - col("l_discount"))) +df = df.with_columns( + l_year=F.datepart(lit("year"), col("l_shipdate")).cast(pa.int32()), + volume=col("l_extendedprice") * (lit(1.0) - col("l_discount")), +) -# Aggregate the results +# Aggregate and sort per the spec. df = df.aggregate( - [col("supp_nation"), col("cust_nation"), col("l_year")], + ["supp_nation", "cust_nation", "l_year"], [F.sum(col("volume")).alias("revenue")], -) - -# Sort based on problem statement requirements -df = df.sort(col("supp_nation").sort(), col("cust_nation").sort(), col("l_year").sort()) +).sort_by("supp_nation", "cust_nation", "l_year") df.show() diff --git a/examples/tpch/q08_market_share.py b/examples/tpch/q08_market_share.py index 4bf50efba..dd7bacedb 100644 --- a/examples/tpch/q08_market_share.py +++ b/examples/tpch/q08_market_share.py @@ -25,24 +25,61 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + o_year, + sum(case + when nation = 'BRAZIL' then volume + else 0 + end) / sum(volume) as mkt_share + from + ( + select + extract(year from o_orderdate) as o_year, + l_extendedprice * (1 - l_discount) as volume, + n2.n_name as nation + from + part, + supplier, + lineitem, + orders, + customer, + nation n1, + nation n2, + region + where + p_partkey = l_partkey + and s_suppkey = l_suppkey + and l_orderkey = o_orderkey + and o_custkey = c_custkey + and c_nationkey = n1.n_nationkey + and n1.n_regionkey = r_regionkey + and r_name = 'AMERICA' + and s_nationkey = n2.n_nationkey + and o_orderdate between date '1995-01-01' and date '1996-12-31' + and p_type = 'ECONOMY ANODIZED STEEL' + ) as all_nations + group by + o_year + order by + o_year; """ -from datetime import datetime +from datetime import date import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path -supplier_nation = lit("BRAZIL") -customer_region = lit("AMERICA") -part_of_interest = lit("ECONOMY ANODIZED STEEL") - -START_DATE = "1995-01-01" -END_DATE = "1996-12-31" +supplier_nation = "BRAZIL" +customer_region = "AMERICA" +part_of_interest = "ECONOMY ANODIZED STEEL" -start_date = lit(datetime.strptime(START_DATE, "%Y-%m-%d").date()) -end_date = lit(datetime.strptime(END_DATE, "%Y-%m-%d").date()) +START_DATE = date(1995, 1, 1) +END_DATE = date(1996, 12, 31) # Load the dataframes we need @@ -74,105 +111,57 @@ # Limit orders to those in the specified range -df_orders = df_orders.filter(col("o_orderdate") >= start_date).filter( - col("o_orderdate") <= end_date -) - -# Part 1: Find customers in the region - -# We want customers in region specified by region_of_interest. This will be used to compute -# the total sales of the part of interest. We want to know of those sales what fraction -# was supplied by the nation of interest. There is no guarantee that the nation of -# interest is within the region of interest. - -# First we find all the sales that make up the basis. - -df_regional_customers = df_region.filter(col("r_name") == customer_region) - -# After this join we have all of the possible sales nations -df_regional_customers = df_regional_customers.join( - df_nation, left_on=["r_regionkey"], right_on=["n_regionkey"], how="inner" -) - -# Now find the possible customers -df_regional_customers = df_regional_customers.join( - df_customer, left_on=["n_nationkey"], right_on=["c_nationkey"], how="inner" -) - -# Next find orders for these customers -df_regional_customers = df_regional_customers.join( - df_orders, left_on=["c_custkey"], right_on=["o_custkey"], how="inner" -) - -# Find all line items from these orders -df_regional_customers = df_regional_customers.join( - df_lineitem, left_on=["o_orderkey"], right_on=["l_orderkey"], how="inner" -) - -# Limit to the part of interest -df_regional_customers = df_regional_customers.join( - df_part, left_on=["l_partkey"], right_on=["p_partkey"], how="inner" -) - -# Compute the volume for each line item -df_regional_customers = df_regional_customers.with_column( - "volume", col("l_extendedprice") * (lit(1.0) - col("l_discount")) -) - -# Part 2: Find suppliers from the nation - -# Now that we have all of the sales of that part in the specified region, we need -# to determine which of those came from suppliers in the nation we are interested in. - -df_national_suppliers = df_nation.filter(col("n_name") == supplier_nation) - -# Determine the suppliers by the limited nation key we have in our single row df above -df_national_suppliers = df_national_suppliers.join( - df_supplier, left_on=["n_nationkey"], right_on=["s_nationkey"], how="inner" -) - -# When we join to the customer dataframe, we don't want to confuse other columns, so only -# select the supplier key that we need -df_national_suppliers = df_national_suppliers.select("s_suppkey") - - -# Part 3: Combine suppliers and customers and compute the market share - -# Now we can do a left outer join on the suppkey. Those line items from other suppliers -# will get a null value. We can check for the existence of this null to compute a volume -# column only from suppliers in the nation we are evaluating. - -df = df_regional_customers.join( - df_national_suppliers, left_on=["l_suppkey"], right_on=["s_suppkey"], how="left" -) - -# Use a case statement to compute the volume sold by suppliers in the nation of interest -df = df.with_column( - "national_volume", - F.case(col("s_suppkey").is_null()) - .when(lit(value=False), col("volume")) - .otherwise(lit(0.0)), -) - -df = df.with_column( - "o_year", F.datepart(lit("year"), col("o_orderdate")).cast(pa.int32()) -) - - -# Lastly, sum up the results - -df = df.aggregate( - [col("o_year")], - [ - F.sum(col("volume")).alias("volume"), - F.sum(col("national_volume")).alias("national_volume"), - ], +df_orders = df_orders.filter( + col("o_orderdate") >= lit(START_DATE), col("o_orderdate") <= lit(END_DATE) +) + +# Pair each supplier with its nation name so every regional-customer row +# below carries the supplier's nation and can be filtered inside the +# aggregate with ``F.sum(..., filter=...)``. + +df_supplier_with_nation = df_supplier.join( + df_nation, left_on="s_nationkey", right_on="n_nationkey" +).select("s_suppkey", col("n_name").alias("supp_nation")) + +# Build every (part, lineitem, order, customer) row for customers in the +# target region ordering the target part. Each row carries the supplier's +# nation so we can aggregate on it below. + +df = ( + df_region.filter(col("r_name") == customer_region) + .join(df_nation, left_on="r_regionkey", right_on="n_regionkey") + .join(df_customer, left_on="n_nationkey", right_on="c_nationkey") + .join(df_orders, left_on="c_custkey", right_on="o_custkey") + .join(df_lineitem, left_on="o_orderkey", right_on="l_orderkey") + .join(df_part, left_on="l_partkey", right_on="p_partkey") + .join(df_supplier_with_nation, left_on="l_suppkey", right_on="s_suppkey") + .with_columns( + volume=col("l_extendedprice") * (lit(1.0) - col("l_discount")), + o_year=F.datepart(lit("year"), col("o_orderdate")).cast(pa.int32()), + ) +) + +# Aggregate the total and national volumes per year via the ``filter`` +# kwarg on ``F.sum`` (DataFrame form of SQL ``sum(... ) FILTER (WHERE ...)``). +# ``coalesce`` handles the case where no sale came from the target nation +# for a given year. +df = ( + df.aggregate( + ["o_year"], + [ + F.sum(col("volume"), filter=col("supp_nation") == supplier_nation).alias( + "national_volume" + ), + F.sum(col("volume")).alias("total_volume"), + ], + ) + .select( + "o_year", + (F.coalesce(col("national_volume"), lit(0.0)) / col("total_volume")).alias( + "mkt_share" + ), + ) + .sort_by("o_year") ) -df = df.select( - col("o_year"), (F.col("national_volume") / F.col("volume")).alias("mkt_share") -) - -df = df.sort(col("o_year").sort()) - df.show() diff --git a/examples/tpch/q09_product_type_profit_measure.py b/examples/tpch/q09_product_type_profit_measure.py index e2abbd095..ec68a2ab7 100644 --- a/examples/tpch/q09_product_type_profit_measure.py +++ b/examples/tpch/q09_product_type_profit_measure.py @@ -27,6 +27,41 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + nation, + o_year, + sum(amount) as sum_profit + from + ( + select + n_name as nation, + extract(year from o_orderdate) as o_year, + l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity as amount + from + part, + supplier, + lineitem, + partsupp, + orders, + nation + where + s_suppkey = l_suppkey + and ps_suppkey = l_suppkey + and ps_partkey = l_partkey + and p_partkey = l_partkey + and o_orderkey = l_orderkey + and s_nationkey = n_nationkey + and p_name like '%green%' + ) as profit + group by + nation, + o_year + order by + nation, + o_year desc; """ import pyarrow as pa @@ -34,7 +69,7 @@ from datafusion import functions as F from util import get_data_path -part_color = lit("green") +part_color = "green" # Load the dataframes we need @@ -62,37 +97,35 @@ "n_nationkey", "n_name", "n_regionkey" ) -# Limit possible parts to the color specified -df = df_part.filter(F.strpos(col("p_name"), part_color) > lit(0)) - -# We have a series of joins that get us to limit down to the line items we need -df = df.join(df_lineitem, left_on=["p_partkey"], right_on=["l_partkey"], how="inner") -df = df.join(df_supplier, left_on=["l_suppkey"], right_on=["s_suppkey"], how="inner") -df = df.join(df_orders, left_on=["l_orderkey"], right_on=["o_orderkey"], how="inner") -df = df.join( - df_partsupp, - left_on=["l_suppkey", "l_partkey"], - right_on=["ps_suppkey", "ps_partkey"], - how="inner", +# Limit possible parts to the color specified, then walk the joins down to the +# line-item rows we need and attach the supplier's nation. ``F.contains`` +# maps directly to the reference SQL's ``p_name like '%green%'``. +df = ( + df_part.filter(F.contains(col("p_name"), lit(part_color))) + .join(df_lineitem, left_on="p_partkey", right_on="l_partkey") + .join(df_supplier, left_on="l_suppkey", right_on="s_suppkey") + .join(df_orders, left_on="l_orderkey", right_on="o_orderkey") + .join( + df_partsupp, + left_on=["l_suppkey", "l_partkey"], + right_on=["ps_suppkey", "ps_partkey"], + ) + .join(df_nation, left_on="s_nationkey", right_on="n_nationkey") ) -df = df.join(df_nation, left_on=["s_nationkey"], right_on=["n_nationkey"], how="inner") # Compute the intermediate values and limit down to the expressions we need df = df.select( col("n_name").alias("nation"), F.datepart(lit("year"), col("o_orderdate")).cast(pa.int32()).alias("o_year"), ( - (col("l_extendedprice") * (lit(1) - col("l_discount"))) - - (col("ps_supplycost") * col("l_quantity")) + col("l_extendedprice") * (lit(1) - col("l_discount")) + - col("ps_supplycost") * col("l_quantity") ).alias("amount"), ) -# Sum up the values by nation and year -df = df.aggregate( - [col("nation"), col("o_year")], [F.sum(col("amount")).alias("profit")] +# Sum up the values by nation and year, then sort per the spec. +df = df.aggregate(["nation", "o_year"], [F.sum(col("amount")).alias("profit")]).sort( + "nation", col("o_year").sort(ascending=False) ) -# Sort according to the problem specification -df = df.sort(col("nation").sort(), col("o_year").sort(ascending=False)) - df.show() diff --git a/examples/tpch/q10_returned_item_reporting.py b/examples/tpch/q10_returned_item_reporting.py index ed822e264..e6532517e 100644 --- a/examples/tpch/q10_returned_item_reporting.py +++ b/examples/tpch/q10_returned_item_reporting.py @@ -27,20 +27,50 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + c_custkey, + c_name, + sum(l_extendedprice * (1 - l_discount)) as revenue, + c_acctbal, + n_name, + c_address, + c_phone, + c_comment + from + customer, + orders, + lineitem, + nation + where + c_custkey = o_custkey + and l_orderkey = o_orderkey + and o_orderdate >= date '1993-10-01' + and o_orderdate < date '1993-10-01' + interval '3' month + and l_returnflag = 'R' + and c_nationkey = n_nationkey + group by + c_custkey, + c_name, + c_acctbal, + c_phone, + n_name, + c_address, + c_comment + order by + revenue desc limit 20; """ -from datetime import datetime +from datetime import date -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path -DATE_START_OF_QUARTER = "1993-10-01" - -date_start_of_quarter = lit(datetime.strptime(DATE_START_OF_QUARTER, "%Y-%m-%d").date()) - -interval_one_quarter = lit(pa.scalar((0, 92, 0), type=pa.month_day_nano_interval())) +QUARTER_START = date(1993, 10, 1) +QUARTER_END = date(1994, 1, 1) # Load the dataframes we need @@ -66,44 +96,40 @@ ) # limit to returns -df_lineitem = df_lineitem.filter(col("l_returnflag") == lit("R")) +df_lineitem = df_lineitem.filter(col("l_returnflag") == "R") # Rather than aggregate by all of the customer fields as you might do looking at the specification, # we can aggregate by o_custkey and then join in the customer data at the end. -df = df_orders.filter(col("o_orderdate") >= date_start_of_quarter).filter( - col("o_orderdate") < date_start_of_quarter + interval_one_quarter +df = ( + df_orders.filter( + col("o_orderdate") >= lit(QUARTER_START), + col("o_orderdate") < lit(QUARTER_END), + ) + .join(df_lineitem, left_on="o_orderkey", right_on="l_orderkey") + .aggregate( + ["o_custkey"], + [F.sum(col("l_extendedprice") * (lit(1) - col("l_discount"))).alias("revenue")], + ) ) -df = df.join(df_lineitem, left_on=["o_orderkey"], right_on=["l_orderkey"], how="inner") - -# Compute the revenue -df = df.aggregate( - [col("o_custkey")], - [F.sum(col("l_extendedprice") * (lit(1) - col("l_discount"))).alias("revenue")], +# Now join in the customer data, project the spec's output columns, and take the top 20. +df = ( + df.join(df_customer, left_on="o_custkey", right_on="c_custkey") + .join(df_nation, left_on="c_nationkey", right_on="n_nationkey") + .select( + "c_custkey", + "c_name", + "revenue", + "c_acctbal", + "n_name", + "c_address", + "c_phone", + "c_comment", + ) + .sort(col("revenue").sort(ascending=False)) + .limit(20) ) -# Now join in the customer data -df = df.join(df_customer, left_on=["o_custkey"], right_on=["c_custkey"], how="inner") -df = df.join(df_nation, left_on=["c_nationkey"], right_on=["n_nationkey"], how="inner") - -# These are the columns the problem statement requires -df = df.select( - "c_custkey", - "c_name", - "revenue", - "c_acctbal", - "n_name", - "c_address", - "c_phone", - "c_comment", -) - -# Sort the results in descending order -df = df.sort(col("revenue").sort(ascending=False)) - -# Only return the top 20 results -df = df.limit(20) - df.show() diff --git a/examples/tpch/q11_important_stock_identification.py b/examples/tpch/q11_important_stock_identification.py index de309fa64..1f40bbdad 100644 --- a/examples/tpch/q11_important_stock_identification.py +++ b/examples/tpch/q11_important_stock_identification.py @@ -25,6 +25,36 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + ps_partkey, + sum(ps_supplycost * ps_availqty) as value + from + partsupp, + supplier, + nation + where + ps_suppkey = s_suppkey + and s_nationkey = n_nationkey + and n_name = 'GERMANY' + group by + ps_partkey having + sum(ps_supplycost * ps_availqty) > ( + select + sum(ps_supplycost * ps_availqty) * 0.0001000000 + from + partsupp, + supplier, + nation + where + ps_suppkey = s_suppkey + and s_nationkey = n_nationkey + and n_name = 'GERMANY' + ) + order by + value desc; """ from datafusion import SessionContext, WindowFrame, col, lit @@ -49,39 +79,30 @@ "n_nationkey", "n_name" ) -# limit to returns -df_nation = df_nation.filter(col("n_name") == lit(NATION)) - -# Find part supplies of within this target nation - -df = df_nation.join( - df_supplier, left_on=["n_nationkey"], right_on=["s_nationkey"], how="inner" +# Restrict to the target nation, then walk to partsupp rows via the supplier +# join. Aggregate the per-part inventory value. +df = ( + df_nation.filter(col("n_name") == NATION) + .join(df_supplier, left_on="n_nationkey", right_on="s_nationkey") + .join(df_partsupp, left_on="s_suppkey", right_on="ps_suppkey") + .with_column("value", col("ps_supplycost") * col("ps_availqty")) + .aggregate(["ps_partkey"], [F.sum(col("value")).alias("value")]) ) -df = df.join(df_partsupp, left_on=["s_suppkey"], right_on=["ps_suppkey"], how="inner") - - -# Compute the value of individual parts -df = df.with_column("value", col("ps_supplycost") * col("ps_availqty")) - -# Compute total value of specific parts -df = df.aggregate([col("ps_partkey")], [F.sum(col("value")).alias("value")]) - -# By default window functions go from unbounded preceding to current row, but we want -# to compute this sum across all rows -window_frame = WindowFrame("rows", None, None) - -df = df.with_column( - "total_value", F.sum(col("value")).over(Window(window_frame=window_frame)) +# A window function evaluated over the entire output produces a scalar grand +# total that can be referenced row-by-row in the filter — a DataFrame-native +# stand-in for the SQL HAVING ... > (SELECT SUM(...) * FRACTION ...) pattern. +# The default frame is "UNBOUNDED PRECEDING to CURRENT ROW"; override to the +# full partition for the grand total. +whole_frame = WindowFrame("rows", None, None) + +df = ( + df.with_column( + "total_value", F.sum(col("value")).over(Window(window_frame=whole_frame)) + ) + .filter(col("value") / col("total_value") >= lit(FRACTION)) + .select("ps_partkey", "value") + .sort(col("value").sort(ascending=False)) ) -# Limit to the parts for which there is a significant value based on the fraction of the total -df = df.filter(col("value") / col("total_value") >= lit(FRACTION)) - -# We only need to report on these two columns -df = df.select("ps_partkey", "value") - -# Sort in descending order of value -df = df.sort(col("value").sort(ascending=False)) - df.show() diff --git a/examples/tpch/q12_ship_mode_order_priority.py b/examples/tpch/q12_ship_mode_order_priority.py index 9071597f0..fb78fe3c2 100644 --- a/examples/tpch/q12_ship_mode_order_priority.py +++ b/examples/tpch/q12_ship_mode_order_priority.py @@ -27,18 +27,49 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + l_shipmode, + sum(case + when o_orderpriority = '1-URGENT' + or o_orderpriority = '2-HIGH' + then 1 + else 0 + end) as high_line_count, + sum(case + when o_orderpriority <> '1-URGENT' + and o_orderpriority <> '2-HIGH' + then 1 + else 0 + end) as low_line_count + from + orders, + lineitem + where + o_orderkey = l_orderkey + and l_shipmode in ('MAIL', 'SHIP') + and l_commitdate < l_receiptdate + and l_shipdate < l_commitdate + and l_receiptdate >= date '1994-01-01' + and l_receiptdate < date '1994-01-01' + interval '1' year + group by + l_shipmode + order by + l_shipmode; """ -from datetime import datetime +from datetime import date -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path SHIP_MODE_1 = "MAIL" SHIP_MODE_2 = "SHIP" -DATE_OF_INTEREST = "1994-01-01" +YEAR_START = date(1994, 1, 1) +YEAR_END = date(1995, 1, 1) # Load the dataframes we need @@ -51,63 +82,30 @@ "l_orderkey", "l_shipmode", "l_commitdate", "l_shipdate", "l_receiptdate" ) -date = datetime.strptime(DATE_OF_INTEREST, "%Y-%m-%d").date() - -interval = pa.scalar((0, 365, 0), type=pa.month_day_nano_interval()) - - -df = df_lineitem.filter(col("l_receiptdate") >= lit(date)).filter( - col("l_receiptdate") < lit(date) + lit(interval) -) - -# Note: It is not recommended to use array_has because it treats the second argument as an argument -# so if you pass it col("l_shipmode") it will pass the entire array to process which is very slow. -# Instead check the position of the entry is not null. -df = df.filter( - ~F.array_position( - F.make_array(lit(SHIP_MODE_1), lit(SHIP_MODE_2)), col("l_shipmode") - ).is_null() -) - -# Since we have only two values, it's much easier to do this as a filter where the l_shipmode -# matches either of the two values, but we want to show doing some array operations in this -# example. If you want to see this done with filters, comment out the above line and uncomment -# this one. -# df = df.filter((col("l_shipmode") == lit(SHIP_MODE_1)) | (col("l_shipmode") == lit(SHIP_MODE_2))) # noqa: ERA001 +df = df_lineitem.filter( + col("l_receiptdate") >= lit(YEAR_START), + col("l_receiptdate") < lit(YEAR_END), + # ``in_list`` maps directly to ``l_shipmode in (...)`` from the SQL. + F.in_list(col("l_shipmode"), [lit(SHIP_MODE_1), lit(SHIP_MODE_2)]), + col("l_shipdate") < col("l_commitdate"), + col("l_commitdate") < col("l_receiptdate"), +).join(df_orders, left_on="l_orderkey", right_on="o_orderkey") -# We need order priority, so join order df to line item -df = df.join(df_orders, left_on=["l_orderkey"], right_on=["o_orderkey"], how="inner") +# Flag each line item as belonging to a high-priority order or not. +high_priorities = [lit("1-URGENT"), lit("2-HIGH")] +is_high = F.in_list(col("o_orderpriority"), high_priorities) +is_low = F.in_list(col("o_orderpriority"), high_priorities, negated=True) -# Restrict to line items we care about based on the problem statement. -df = df.filter(col("l_commitdate") < col("l_receiptdate")) - -df = df.filter(col("l_shipdate") < col("l_commitdate")) - -df = df.with_column( - "high_line_value", - F.case(col("o_orderpriority")) - .when(lit("1-URGENT"), lit(1)) - .when(lit("2-HIGH"), lit(1)) - .otherwise(lit(0)), -) - -# Aggregate the results +# Count the high-priority and low-priority lineitems per ship mode via the +# ``filter`` kwarg on ``F.count`` (DataFrame form of SQL's ``count(*) +# FILTER (WHERE ...)``). df = df.aggregate( - [col("l_shipmode")], + ["l_shipmode"], [ - F.sum(col("high_line_value")).alias("high_line_count"), - F.count(col("high_line_value")).alias("all_lines_count"), + F.count(col("o_orderkey"), filter=is_high).alias("high_line_count"), + F.count(col("o_orderkey"), filter=is_low).alias("low_line_count"), ], -) - -# Compute the final output -df = df.select( - col("l_shipmode"), - col("high_line_count"), - (col("all_lines_count") - col("high_line_count")).alias("low_line_count"), -) - -df = df.sort(col("l_shipmode").sort()) +).sort_by("l_shipmode") df.show() diff --git a/examples/tpch/q13_customer_distribution.py b/examples/tpch/q13_customer_distribution.py index 93f082ea3..37c0b93f6 100644 --- a/examples/tpch/q13_customer_distribution.py +++ b/examples/tpch/q13_customer_distribution.py @@ -26,6 +26,29 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + c_count, + count(*) as custdist + from + ( + select + c_custkey, + count(o_orderkey) + from + customer left outer join orders on + c_custkey = o_custkey + and o_comment not like '%special%requests%' + group by + c_custkey + ) as c_orders (c_custkey, c_count) + group by + c_count + order by + custdist desc, + c_count desc; """ from datafusion import SessionContext, col, lit @@ -49,20 +72,16 @@ F.regexp_match(col("o_comment"), lit(f"{WORD_1}.?*{WORD_2}")).is_null() ) -# Since we may have customers with no orders we must do a left join -df = df_customer.join( - df_orders, left_on=["c_custkey"], right_on=["o_custkey"], how="left" -) - -# Find the number of orders for each customer -df = df.aggregate([col("c_custkey")], [F.count(col("o_custkey")).alias("c_count")]) - -# Ultimately we want to know the number of customers that have that customer count -df = df.aggregate([col("c_count")], [F.count(col("c_count")).alias("custdist")]) - -# We want to order the results by the highest number of customers per count -df = df.sort( - col("custdist").sort(ascending=False), col("c_count").sort(ascending=False) +# Customers with no orders still participate, so this is a left join. Count the +# orders per customer, then count customers per order-count value. +df = ( + df_customer.join(df_orders, left_on="c_custkey", right_on="o_custkey", how="left") + .aggregate(["c_custkey"], [F.count(col("o_custkey")).alias("c_count")]) + .aggregate(["c_count"], [F.count_star().alias("custdist")]) + .sort( + col("custdist").sort(ascending=False), + col("c_count").sort(ascending=False), + ) ) df.show() diff --git a/examples/tpch/q14_promotion_effect.py b/examples/tpch/q14_promotion_effect.py index d62f76e3c..08f4f054d 100644 --- a/examples/tpch/q14_promotion_effect.py +++ b/examples/tpch/q14_promotion_effect.py @@ -24,20 +24,32 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + 100.00 * sum(case + when p_type like 'PROMO%' + then l_extendedprice * (1 - l_discount) + else 0 + end) / sum(l_extendedprice * (1 - l_discount)) as promo_revenue + from + lineitem, + part + where + l_partkey = p_partkey + and l_shipdate >= date '1995-09-01' + and l_shipdate < date '1995-09-01' + interval '1' month; """ -from datetime import datetime +from datetime import date -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path -DATE = "1995-09-01" - -date_of_interest = lit(datetime.strptime(DATE, "%Y-%m-%d").date()) - -interval_one_month = lit(pa.scalar((0, 30, 0), type=pa.month_day_nano_interval())) +MONTH_START = date(1995, 9, 1) +MONTH_END = date(1995, 10, 1) # Load the dataframes we need @@ -49,37 +61,30 @@ df_part = ctx.read_parquet(get_data_path("part.parquet")).select("p_partkey", "p_type") -# Check part type begins with PROMO -df_part = df_part.filter( - F.substring(col("p_type"), lit(0), lit(6)) == lit("PROMO") -).with_column("promo_factor", lit(1.0)) - -df_lineitem = df_lineitem.filter(col("l_shipdate") >= date_of_interest).filter( - col("l_shipdate") < date_of_interest + interval_one_month -) - -# Left join so we can sum up the promo parts different from other parts -df = df_lineitem.join( - df_part, left_on=["l_partkey"], right_on=["p_partkey"], how="left" -) - -# Make a factor of 1.0 if it is a promotion, 0.0 otherwise -df = df.with_column("promo_factor", F.coalesce(col("promo_factor"), lit(0.0))) -df = df.with_column("revenue", col("l_extendedprice") * (lit(1.0) - col("l_discount"))) - - -# Sum up the promo and total revenue -df = df.aggregate( - [], - [ - F.sum(col("promo_factor") * col("revenue")).alias("promo_revenue"), - F.sum(col("revenue")).alias("total_revenue"), - ], -) - -# Return the percentage of revenue from promotions -df = df.select( - (lit(100.0) * col("promo_revenue") / col("total_revenue")).alias("promo_revenue") +# Restrict the line items to the month of interest, join the matching part +# rows, and aggregate revenue totals with a ``filter`` clause on the promo +# sum — the DataFrame form of SQL ``sum(... ) FILTER (WHERE ...)``. +revenue = col("l_extendedprice") * (lit(1.0) - col("l_discount")) +is_promo = F.starts_with(col("p_type"), lit("PROMO")) + +df = ( + df_lineitem.filter( + col("l_shipdate") >= lit(MONTH_START), + col("l_shipdate") < lit(MONTH_END), + ) + .join(df_part, left_on="l_partkey", right_on="p_partkey") + .aggregate( + [], + [ + F.sum(revenue, filter=is_promo).alias("promo_revenue"), + F.sum(revenue).alias("total_revenue"), + ], + ) + .select( + (lit(100.0) * col("promo_revenue") / col("total_revenue")).alias( + "promo_revenue" + ) + ) ) df.show() diff --git a/examples/tpch/q15_top_supplier.py b/examples/tpch/q15_top_supplier.py index 5128937a7..01c38b9f8 100644 --- a/examples/tpch/q15_top_supplier.py +++ b/examples/tpch/q15_top_supplier.py @@ -24,21 +24,50 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + create view revenue0 (supplier_no, total_revenue) as + select + l_suppkey, + sum(l_extendedprice * (1 - l_discount)) + from + lineitem + where + l_shipdate >= date '1996-01-01' + and l_shipdate < date '1996-01-01' + interval '3' month + group by + l_suppkey; + select + s_suppkey, + s_name, + s_address, + s_phone, + total_revenue + from + supplier, + revenue0 + where + s_suppkey = supplier_no + and total_revenue = ( + select + max(total_revenue) + from + revenue0 + ) + order by + s_suppkey; + drop view revenue0; """ -from datetime import datetime +from datetime import date -import pyarrow as pa -from datafusion import SessionContext, WindowFrame, col, lit +from datafusion import SessionContext, col, lit from datafusion import functions as F -from datafusion.expr import Window from util import get_data_path -DATE = "1996-01-01" - -date_of_interest = lit(datetime.strptime(DATE, "%Y-%m-%d").date()) - -interval_3_months = lit(pa.scalar((0, 91, 0), type=pa.month_day_nano_interval())) +QUARTER_START = date(1996, 1, 1) +QUARTER_END = date(1996, 4, 1) # Load the dataframes we need @@ -54,38 +83,29 @@ "s_phone", ) -# Limit line items to the quarter of interest -df_lineitem = df_lineitem.filter(col("l_shipdate") >= date_of_interest).filter( - col("l_shipdate") < date_of_interest + interval_3_months -) +# Per-supplier revenue over the quarter of interest. +revenue = col("l_extendedprice") * (lit(1) - col("l_discount")) -df = df_lineitem.aggregate( - [col("l_suppkey")], - [ - F.sum(col("l_extendedprice") * (lit(1) - col("l_discount"))).alias( - "total_revenue" - ) - ], -) +per_supplier_revenue = df_lineitem.filter( + col("l_shipdate") >= lit(QUARTER_START), + col("l_shipdate") < lit(QUARTER_END), +).aggregate(["l_suppkey"], [F.sum(revenue).alias("total_revenue")]) -# Use a window function to find the maximum revenue across the entire dataframe -window_frame = WindowFrame("rows", None, None) -df = df.with_column( - "max_revenue", - F.max(col("total_revenue")).over(Window(window_frame=window_frame)), +# Compute the grand maximum revenue separately and join on equality — the +# DataFrame stand-in for the reference SQL's +# ``total_revenue = (select max(total_revenue) from revenue0)`` subquery. +max_revenue = per_supplier_revenue.aggregate( + [], [F.max(col("total_revenue")).alias("max_rev")] ) -# Find all suppliers whose total revenue is the same as the maximum -df = df.filter(col("total_revenue") == col("max_revenue")) - -# Now that we know the supplier(s) with maximum revenue, get the rest of their information -# from the supplier table -df = df.join(df_supplier, left_on=["l_suppkey"], right_on=["s_suppkey"], how="inner") +top_suppliers = per_supplier_revenue.join_on( + max_revenue, col("total_revenue") == col("max_rev") +).select("l_suppkey", "total_revenue") -# Return only the columns requested -df = df.select("s_suppkey", "s_name", "s_address", "s_phone", "total_revenue") - -# If we have more than one, sort by supplier number (suppkey) -df = df.sort(col("s_suppkey").sort()) +df = ( + df_supplier.join(top_suppliers, left_on="s_suppkey", right_on="l_suppkey") + .select("s_suppkey", "s_name", "s_address", "s_phone", "total_revenue") + .sort_by("s_suppkey") +) df.show() diff --git a/examples/tpch/q16_part_supplier_relationship.py b/examples/tpch/q16_part_supplier_relationship.py index 65043ffda..ddeadff5f 100644 --- a/examples/tpch/q16_part_supplier_relationship.py +++ b/examples/tpch/q16_part_supplier_relationship.py @@ -26,9 +26,41 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + p_brand, + p_type, + p_size, + count(distinct ps_suppkey) as supplier_cnt + from + partsupp, + part + where + p_partkey = ps_partkey + and p_brand <> 'Brand#45' + and p_type not like 'MEDIUM POLISHED%' + and p_size in (49, 14, 23, 45, 19, 3, 36, 9) + and ps_suppkey not in ( + select + s_suppkey + from + supplier + where + s_comment like '%Customer%Complaints%' + ) + group by + p_brand, + p_type, + p_size + order by + supplier_cnt desc, + p_brand, + p_type, + p_size; """ -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path @@ -52,39 +84,36 @@ ) df_unwanted_suppliers = df_supplier.filter( - ~F.regexp_match(col("s_comment"), lit("Customer.?*Complaints")).is_null() + F.regexp_like(col("s_comment"), lit("Customer.*Complaints")) ) -# Remove unwanted suppliers +# Remove unwanted suppliers via an anti join (DataFrame form of NOT IN). df_partsupp = df_partsupp.join( - df_unwanted_suppliers, left_on=["ps_suppkey"], right_on=["s_suppkey"], how="anti" + df_unwanted_suppliers, left_on="ps_suppkey", right_on="s_suppkey", how="anti" ) -# Select the parts we are interested in -df_part = df_part.filter(col("p_brand") != lit(BRAND)) +# Select the parts we are interested in. df_part = df_part.filter( - F.substring(col("p_type"), lit(0), lit(len(TYPE_TO_IGNORE) + 1)) - != lit(TYPE_TO_IGNORE) -) - -# Python conversion of integer to literal casts it to int64 but the data for -# part size is stored as an int32, so perform a cast. Then check to find if the part -# size is within the array of possible sizes by checking the position of it is not -# null. -p_sizes = F.make_array(*[lit(s).cast(pa.int32()) for s in SIZES_OF_INTEREST]) -df_part = df_part.filter(~F.array_position(p_sizes, col("p_size")).is_null()) - -df = df_part.join( - df_partsupp, left_on=["p_partkey"], right_on=["ps_partkey"], how="inner" + col("p_brand") != BRAND, + ~F.starts_with(col("p_type"), lit(TYPE_TO_IGNORE)), + F.in_list(col("p_size"), [lit(s) for s in SIZES_OF_INTEREST]), ) -df = df.select("p_brand", "p_type", "p_size", "ps_suppkey").distinct() - -df = df.aggregate( - [col("p_brand"), col("p_type"), col("p_size")], - [F.count(col("ps_suppkey")).alias("supplier_cnt")], +# For each (brand, type, size), count the distinct suppliers remaining. +df = ( + df_part.join(df_partsupp, left_on="p_partkey", right_on="ps_partkey") + .select("p_brand", "p_type", "p_size", "ps_suppkey") + .distinct() + .aggregate( + ["p_brand", "p_type", "p_size"], + [F.count(col("ps_suppkey")).alias("supplier_cnt")], + ) + .sort( + col("supplier_cnt").sort(ascending=False), + "p_brand", + "p_type", + "p_size", + ) ) -df = df.sort(col("supplier_cnt").sort(ascending=False)) - df.show() diff --git a/examples/tpch/q17_small_quantity_order.py b/examples/tpch/q17_small_quantity_order.py index 5ccb38422..f2229171f 100644 --- a/examples/tpch/q17_small_quantity_order.py +++ b/examples/tpch/q17_small_quantity_order.py @@ -26,6 +26,26 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + sum(l_extendedprice) / 7.0 as avg_yearly + from + lineitem, + part + where + p_partkey = l_partkey + and p_brand = 'Brand#23' + and p_container = 'MED BOX' + and l_quantity < ( + select + 0.2 * avg(l_quantity) + from + lineitem + where + l_partkey = p_partkey + ); """ from datafusion import SessionContext, WindowFrame, col, lit @@ -47,29 +67,23 @@ "l_partkey", "l_quantity", "l_extendedprice" ) -# Limit to the problem statement's brand and container types -df = df_part.filter(col("p_brand") == lit(BRAND)).filter( - col("p_container") == lit(CONTAINER) -) - -# Combine data -df = df.join(df_lineitem, left_on=["p_partkey"], right_on=["l_partkey"], how="inner") +# Limit to parts of the target brand/container, join their line items, and +# attach the per-part average quantity via a partitioned window function — +# the DataFrame form of the SQL's correlated ``avg(l_quantity)`` subquery. +whole_frame = WindowFrame("rows", None, None) -# Find the average quantity -window_frame = WindowFrame("rows", None, None) -df = df.with_column( - "avg_quantity", - F.avg(col("l_quantity")).over( - Window(partition_by=[col("l_partkey")], window_frame=window_frame) - ), +df = ( + df_part.filter(col("p_brand") == BRAND, col("p_container") == CONTAINER) + .join(df_lineitem, left_on="p_partkey", right_on="l_partkey") + .with_column( + "avg_quantity", + F.avg(col("l_quantity")).over( + Window(partition_by=[col("l_partkey")], window_frame=whole_frame) + ), + ) + .filter(col("l_quantity") < lit(0.2) * col("avg_quantity")) + .aggregate([], [F.sum(col("l_extendedprice")).alias("total")]) + .select((col("total") / lit(7.0)).alias("avg_yearly")) ) -df = df.filter(col("l_quantity") < lit(0.2) * col("avg_quantity")) - -# Compute the total -df = df.aggregate([], [F.sum(col("l_extendedprice")).alias("total")]) - -# Divide by number of years in the problem statement to get average -df = df.select((col("total") / lit(7)).alias("avg_yearly")) - df.show() diff --git a/examples/tpch/q18_large_volume_customer.py b/examples/tpch/q18_large_volume_customer.py index 834d181c9..23132d60d 100644 --- a/examples/tpch/q18_large_volume_customer.py +++ b/examples/tpch/q18_large_volume_customer.py @@ -24,9 +24,44 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + c_name, + c_custkey, + o_orderkey, + o_orderdate, + o_totalprice, + sum(l_quantity) + from + customer, + orders, + lineitem + where + o_orderkey in ( + select + l_orderkey + from + lineitem + group by + l_orderkey having + sum(l_quantity) > 300 + ) + and c_custkey = o_custkey + and o_orderkey = l_orderkey + group by + c_name, + c_custkey, + o_orderkey, + o_orderdate, + o_totalprice + order by + o_totalprice desc, + o_orderdate limit 100; """ -from datafusion import SessionContext, col, lit +from datafusion import SessionContext, col from datafusion import functions as F from util import get_data_path @@ -46,22 +81,24 @@ "l_orderkey", "l_quantity", "l_extendedprice" ) -df = df_lineitem.aggregate( - [col("l_orderkey")], [F.sum(col("l_quantity")).alias("total_quantity")] +# Find orders whose total quantity exceeds the threshold, then join in the +# order + customer details the problem statement requires and sort. +df = ( + df_lineitem.aggregate( + ["l_orderkey"], [F.sum(col("l_quantity")).alias("total_quantity")] + ) + .filter(col("total_quantity") > QUANTITY) + .join(df_orders, left_on="l_orderkey", right_on="o_orderkey") + .join(df_customer, left_on="o_custkey", right_on="c_custkey") + .select( + "c_name", + "c_custkey", + "o_orderkey", + "o_orderdate", + "o_totalprice", + "total_quantity", + ) + .sort(col("o_totalprice").sort(ascending=False), "o_orderdate") ) -# Limit to orders in which the total quantity is above a threshold -df = df.filter(col("total_quantity") > lit(QUANTITY)) - -# We've identified the orders of interest, now join the additional data -# we are required to report on -df = df.join(df_orders, left_on=["l_orderkey"], right_on=["o_orderkey"], how="inner") -df = df.join(df_customer, left_on=["o_custkey"], right_on=["c_custkey"], how="inner") - -df = df.select( - "c_name", "c_custkey", "o_orderkey", "o_orderdate", "o_totalprice", "total_quantity" -) - -df = df.sort(col("o_totalprice").sort(ascending=False), col("o_orderdate").sort()) - df.show() diff --git a/examples/tpch/q19_discounted_revenue.py b/examples/tpch/q19_discounted_revenue.py index bd492aac0..a2be1c1b7 100644 --- a/examples/tpch/q19_discounted_revenue.py +++ b/examples/tpch/q19_discounted_revenue.py @@ -24,10 +24,47 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + sum(l_extendedprice* (1 - l_discount)) as revenue + from + lineitem, + part + where + ( + p_partkey = l_partkey + and p_brand = 'Brand#12' + and p_container in ('SM CASE', 'SM BOX', 'SM PACK', 'SM PKG') + and l_quantity >= 1 and l_quantity <= 1 + 10 + and p_size between 1 and 5 + and l_shipmode in ('AIR', 'AIR REG') + and l_shipinstruct = 'DELIVER IN PERSON' + ) + or + ( + p_partkey = l_partkey + and p_brand = 'Brand#23' + and p_container in ('MED BAG', 'MED BOX', 'MED PKG', 'MED PACK') + and l_quantity >= 10 and l_quantity <= 10 + 10 + and p_size between 1 and 10 + and l_shipmode in ('AIR', 'AIR REG') + and l_shipinstruct = 'DELIVER IN PERSON' + ) + or + ( + p_partkey = l_partkey + and p_brand = 'Brand#34' + and p_container in ('LG CASE', 'LG BOX', 'LG PACK', 'LG PKG') + and l_quantity >= 20 and l_quantity <= 20 + 10 + and p_size between 1 and 15 + and l_shipmode in ('AIR', 'AIR REG') + and l_shipinstruct = 'DELIVER IN PERSON' + ); """ -import pyarrow as pa -from datafusion import SessionContext, col, lit, udf +from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path @@ -65,72 +102,41 @@ "l_discount", ) -# These limitations apply to all line items, so go ahead and do them first - -df = df_lineitem.filter(col("l_shipinstruct") == lit("DELIVER IN PERSON")) - -df = df.filter( - (col("l_shipmode") == lit("AIR")) | (col("l_shipmode") == lit("AIR REG")) -) +# Filter conditions that apply to every disjunct of the reference SQL's WHERE +# clause — pull them out up front so the per-brand predicate stays focused on +# the brand-specific parts. +df = df_lineitem.filter( + col("l_shipinstruct") == "DELIVER IN PERSON", + F.in_list(col("l_shipmode"), [lit("AIR"), lit("AIR REG")]), +).join(df_part, left_on="l_partkey", right_on="p_partkey") + + +# Build one OR-combined predicate per brand. Each disjunct encodes the +# brand-specific container list, quantity window, and size range from the +# reference SQL. This mirrors the SQL ``where (... brand A ...) or (... brand +# B ...) or (... brand C ...)`` form directly, without a UDF. +def _brand_predicate( + brand: str, min_quantity: int, containers: list[str], max_size: int +): + return ( + (col("p_brand") == brand) + & F.in_list(col("p_container"), [lit(c) for c in containers]) + & col("l_quantity").between(lit(min_quantity), lit(min_quantity + 10)) + & col("p_size").between(lit(1), lit(max_size)) + ) -df = df.join(df_part, left_on=["l_partkey"], right_on=["p_partkey"], how="inner") - - -# Create the user defined function (UDF) definition that does the work -def is_of_interest( - brand_arr: pa.Array, - container_arr: pa.Array, - quantity_arr: pa.Array, - size_arr: pa.Array, -) -> pa.Array: - """ - The purpose of this function is to demonstrate how a UDF works, taking as input a pyarrow Array - and generating a resultant Array. The length of the inputs should match and there should be the - same number of rows in the output. - """ - result = [] - for idx, brand_val in enumerate(brand_arr): - brand = brand_val.as_py() - if brand in items_of_interest: - values_of_interest = items_of_interest[brand] - - container_matches = ( - container_arr[idx].as_py() in values_of_interest["containers"] - ) - - quantity = quantity_arr[idx].as_py() - quantity_matches = ( - values_of_interest["min_quantity"] - <= quantity - <= values_of_interest["min_quantity"] + 10 - ) - - size = size_arr[idx].as_py() - size_matches = 1 <= size <= values_of_interest["max_size"] - - result.append(container_matches and quantity_matches and size_matches) - else: - result.append(False) - - return pa.array(result) - - -# Turn the above function into a UDF that DataFusion can understand -is_of_interest_udf = udf( - is_of_interest, - [pa.utf8(), pa.utf8(), pa.decimal128(15, 2), pa.int32()], - pa.bool_(), - "stable", -) -# Filter results using the above UDF -df = df.filter( - is_of_interest_udf( - col("p_brand"), col("p_container"), col("l_quantity"), col("p_size") +predicate = None +for brand, params in items_of_interest.items(): + part_predicate = _brand_predicate( + brand, + params["min_quantity"], + params["containers"], + params["max_size"], ) -) + predicate = part_predicate if predicate is None else predicate | part_predicate -df = df.aggregate( +df = df.filter(predicate).aggregate( [], [F.sum(col("l_extendedprice") * (lit(1) - col("l_discount"))).alias("revenue")], ) diff --git a/examples/tpch/q20_potential_part_promotion.py b/examples/tpch/q20_potential_part_promotion.py index a25188d31..18f96da97 100644 --- a/examples/tpch/q20_potential_part_promotion.py +++ b/examples/tpch/q20_potential_part_promotion.py @@ -25,17 +25,57 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + s_name, + s_address + from + supplier, + nation + where + s_suppkey in ( + select + ps_suppkey + from + partsupp + where + ps_partkey in ( + select + p_partkey + from + part + where + p_name like 'forest%' + ) + and ps_availqty > ( + select + 0.5 * sum(l_quantity) + from + lineitem + where + l_partkey = ps_partkey + and l_suppkey = ps_suppkey + and l_shipdate >= date '1994-01-01' + and l_shipdate < date '1994-01-01' + interval '1' year + ) + ) + and s_nationkey = n_nationkey + and n_name = 'CANADA' + order by + s_name; """ -from datetime import datetime +from datetime import date -import pyarrow as pa from datafusion import SessionContext, col, lit from datafusion import functions as F from util import get_data_path COLOR_OF_INTEREST = "forest" -DATE_OF_INTEREST = "1994-01-01" +YEAR_START = date(1994, 1, 1) +YEAR_END = date(1995, 1, 1) NATION_OF_INTEREST = "CANADA" # Load the dataframes we need @@ -56,46 +96,48 @@ "n_nationkey", "n_name" ) -date = datetime.strptime(DATE_OF_INTEREST, "%Y-%m-%d").date() - -interval = pa.scalar((0, 365, 0), type=pa.month_day_nano_interval()) - -# Filter down dataframes -df_nation = df_nation.filter(col("n_name") == lit(NATION_OF_INTEREST)) -df_part = df_part.filter( - F.substring(col("p_name"), lit(0), lit(len(COLOR_OF_INTEREST) + 1)) - == lit(COLOR_OF_INTEREST) -) - -df = df_lineitem.filter(col("l_shipdate") >= lit(date)).filter( - col("l_shipdate") < lit(date) + lit(interval) +# Filter down dataframes. ``starts_with`` reads more naturally than an +# explicit substring slice and maps directly to the reference SQL's +# ``p_name like 'forest%'`` clause. +df_nation = df_nation.filter(col("n_name") == NATION_OF_INTEREST) +df_part = df_part.filter(F.starts_with(col("p_name"), lit(COLOR_OF_INTEREST))) + +# Compute the total quantity of interesting parts shipped by each (part, +# supplier) pair within the year of interest. +totals = ( + df_lineitem.filter( + col("l_shipdate") >= lit(YEAR_START), + col("l_shipdate") < lit(YEAR_END), + ) + .join(df_part, left_on="l_partkey", right_on="p_partkey") + .aggregate( + ["l_partkey", "l_suppkey"], + [F.sum(col("l_quantity")).alias("total_sold")], + ) ) -# This will filter down the line items to the parts of interest -df = df.join(df_part, left_on="l_partkey", right_on="p_partkey", how="inner") - -# Compute the total sold and limit ourselves to individual supplier/part combinations -df = df.aggregate( - [col("l_partkey"), col("l_suppkey")], [F.sum(col("l_quantity")).alias("total_sold")] +# Keep only (part, supplier) pairs whose available quantity exceeds 50% of +# the total shipped. The result already contains one row per supplier of +# interest, so we can semi-join the supplier table rather than inner-join +# and deduplicate afterwards. +excess_suppliers = ( + df_partsupp.join( + totals, + left_on=["ps_partkey", "ps_suppkey"], + right_on=["l_partkey", "l_suppkey"], + ) + .filter(col("ps_availqty") > lit(0.5) * col("total_sold")) + .select(col("ps_suppkey").alias("suppkey")) + .distinct() ) -df = df.join( - df_partsupp, - left_on=["l_partkey", "l_suppkey"], - right_on=["ps_partkey", "ps_suppkey"], - how="inner", +# Limit to suppliers in the nation of interest and pick out the two +# requested columns. +df = ( + df_supplier.join(df_nation, left_on="s_nationkey", right_on="n_nationkey") + .join(excess_suppliers, left_on="s_suppkey", right_on="suppkey", how="semi") + .select("s_name", "s_address") + .sort_by("s_name") ) -# Find cases of excess quantity -df.filter(col("ps_availqty") > lit(0.5) * col("total_sold")) - -# We could do these joins earlier, but now limit to the nation of interest suppliers -df = df.join(df_supplier, left_on=["ps_suppkey"], right_on=["s_suppkey"], how="inner") -df = df.join(df_nation, left_on=["s_nationkey"], right_on=["n_nationkey"], how="inner") - -# Restrict to the requested data per the problem statement -df = df.select("s_name", "s_address").distinct() - -df = df.sort(col("s_name").sort()) - df.show() diff --git a/examples/tpch/q21_suppliers_kept_orders_waiting.py b/examples/tpch/q21_suppliers_kept_orders_waiting.py index 4ee9d3733..d98f76ce7 100644 --- a/examples/tpch/q21_suppliers_kept_orders_waiting.py +++ b/examples/tpch/q21_suppliers_kept_orders_waiting.py @@ -24,9 +24,51 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + s_name, + count(*) as numwait + from + supplier, + lineitem l1, + orders, + nation + where + s_suppkey = l1.l_suppkey + and o_orderkey = l1.l_orderkey + and o_orderstatus = 'F' + and l1.l_receiptdate > l1.l_commitdate + and exists ( + select + * + from + lineitem l2 + where + l2.l_orderkey = l1.l_orderkey + and l2.l_suppkey <> l1.l_suppkey + ) + and not exists ( + select + * + from + lineitem l3 + where + l3.l_orderkey = l1.l_orderkey + and l3.l_suppkey <> l1.l_suppkey + and l3.l_receiptdate > l3.l_commitdate + ) + and s_nationkey = n_nationkey + and n_name = 'SAUDI ARABIA' + group by + s_name + order by + numwait desc, + s_name limit 100; """ -from datafusion import SessionContext, col, lit +from datafusion import SessionContext, col from datafusion import functions as F from util import get_data_path @@ -50,65 +92,57 @@ ) # Limit to suppliers in the nation of interest -df_suppliers_of_interest = df_nation.filter(col("n_name") == lit(NATION_OF_INTEREST)) - -df_suppliers_of_interest = df_suppliers_of_interest.join( - df_supplier, left_on="n_nationkey", right_on="s_nationkey", how="inner" +df_suppliers_of_interest = df_nation.filter(col("n_name") == NATION_OF_INTEREST).join( + df_supplier, left_on="n_nationkey", right_on="s_nationkey" ) -# Find the failed orders and all their line items -df = df_orders.filter(col("o_orderstatus") == lit("F")) - -df = df_lineitem.join(df, left_on="l_orderkey", right_on="o_orderkey", how="inner") - -# Identify the line items for which the order is failed due to. -df = df.with_column( - "failed_supp", - F.case(col("l_receiptdate") > col("l_commitdate")) - .when(lit(value=True), col("l_suppkey")) - .end(), +# Line items for orders that have status 'F'. This is the candidate set of +# (order, supplier) pairs we reason about below. +failed_order_lineitems = df_lineitem.join( + df_orders.filter(col("o_orderstatus") == "F"), + left_on="l_orderkey", + right_on="o_orderkey", ) -# There are other ways we could do this but the purpose of this example is to work with rows where -# an element is an array of values. In this case, we will create two columns of arrays. One will be -# an array of all of the suppliers who made up this order. That way we can filter the dataframe for -# only orders where this array is larger than one for multiple supplier orders. The second column -# is all of the suppliers who failed to make their commitment. We can filter the second column for -# arrays with size one. That combination will give us orders that had multiple suppliers where only -# one failed. Use distinct=True in the blow aggregation so we don't get multiple line items from the -# same supplier reported in either array. -df = df.aggregate( - [col("o_orderkey")], - [ - F.array_agg(col("l_suppkey"), distinct=True).alias("all_suppliers"), - F.array_agg( - col("failed_supp"), filter=col("failed_supp").is_not_null(), distinct=True - ).alias("failed_suppliers"), - ], +# Line items whose receipt was late. This corresponds to ``l1`` in the +# reference SQL. +late_lineitems = failed_order_lineitems.filter( + col("l_receiptdate") > col("l_commitdate") ) -# This is the check described above which will identify single failed supplier in a multiple -# supplier order. -df = df.filter(F.array_length(col("failed_suppliers")) == lit(1)).filter( - F.array_length(col("all_suppliers")) > lit(1) +# Orders that had more than one distinct supplier. Expressed as +# ``count(distinct l_suppkey) > 1``. Stands in for the reference SQL's +# ``exists (... l2.l_suppkey <> l1.l_suppkey ...)`` subquery. +multi_supplier_orders = ( + failed_order_lineitems.select("l_orderkey", "l_suppkey") + .distinct() + .aggregate(["l_orderkey"], [F.count_star().alias("n_suppliers")]) + .filter(col("n_suppliers") > 1) + .select("l_orderkey") ) -# Since we have an array we know is exactly one element long, we can extract that single value. -df = df.select( - col("o_orderkey"), F.array_element(col("failed_suppliers"), lit(1)).alias("suppkey") +# Orders where exactly one distinct supplier was late. Stands in for the +# reference SQL's ``not exists (... l3.l_suppkey <> l1.l_suppkey and l3 is +# also late ...)`` subquery: if only one supplier on the order was late, +# nobody else on the same order was late. +single_late_supplier_orders = ( + late_lineitems.select("l_orderkey", "l_suppkey") + .distinct() + .aggregate(["l_orderkey"], [F.count_star().alias("n_late_suppliers")]) + .filter(col("n_late_suppliers") == 1) + .select("l_orderkey") ) -# Join to the supplier of interest list for the nation of interest -df = df.join( - df_suppliers_of_interest, left_on=["suppkey"], right_on=["s_suppkey"], how="inner" +# Keep late line items whose order qualifies on both counts, attach the +# supplier name for suppliers in the nation of interest, count one row per +# qualifying order, and return the top 100. +df = ( + late_lineitems.join(multi_supplier_orders, on="l_orderkey", how="semi") + .join(single_late_supplier_orders, on="l_orderkey", how="semi") + .join(df_suppliers_of_interest, left_on="l_suppkey", right_on="s_suppkey") + .aggregate(["s_name"], [F.count_star().alias("numwait")]) + .sort(col("numwait").sort(ascending=False), "s_name") + .limit(100) ) -# Count how many orders that supplier is the only failed supplier for -df = df.aggregate([col("s_name")], [F.count(col("o_orderkey")).alias("numwait")]) - -# Return in descending order -df = df.sort(col("numwait").sort(ascending=False), col("s_name").sort()) - -df = df.limit(100) - df.show() diff --git a/examples/tpch/q22_global_sales_opportunity.py b/examples/tpch/q22_global_sales_opportunity.py index a2d41b215..5043eeb51 100644 --- a/examples/tpch/q22_global_sales_opportunity.py +++ b/examples/tpch/q22_global_sales_opportunity.py @@ -24,6 +24,46 @@ The above problem statement text is copyrighted by the Transaction Processing Performance Council as part of their TPC Benchmark H Specification revision 2.18.0. + +Reference SQL (from TPC-H specification, used by the benchmark suite):: + + select + cntrycode, + count(*) as numcust, + sum(c_acctbal) as totacctbal + from + ( + select + substring(c_phone from 1 for 2) as cntrycode, + c_acctbal + from + customer + where + substring(c_phone from 1 for 2) in + ('13', '31', '23', '29', '30', '18', '17') + and c_acctbal > ( + select + avg(c_acctbal) + from + customer + where + c_acctbal > 0.00 + and substring(c_phone from 1 for 2) in + ('13', '31', '23', '29', '30', '18', '17') + ) + and not exists ( + select + * + from + orders + where + o_custkey = c_custkey + ) + ) as custsale + group by + cntrycode + order by + cntrycode; """ from datafusion import SessionContext, WindowFrame, col, lit @@ -42,40 +82,36 @@ ) df_orders = ctx.read_parquet(get_data_path("orders.parquet")).select("o_custkey") -# The nation code is a two digit number, but we need to convert it to a string literal -nation_codes = F.make_array(*[lit(str(n)) for n in NATION_CODES]) - -# Use the substring operation to extract the first two characters of the phone number -df = df_customer.with_column("cntrycode", F.substring(col("c_phone"), lit(0), lit(3))) - -# Limit our search to customers with some balance and in the country code above -df = df.filter(col("c_acctbal") > lit(0.0)) -df = df.filter(~F.array_position(nation_codes, col("cntrycode")).is_null()) - -# Compute the average balance. By default, the window frame is from unbounded preceding to the -# current row. We want our frame to cover the entire data frame. -window_frame = WindowFrame("rows", None, None) -df = df.with_column( - "avg_balance", - F.avg(col("c_acctbal")).over(Window(window_frame=window_frame)), -) - -df.show() -# Limit results to customers with above average balance -df = df.filter(col("c_acctbal") > col("avg_balance")) - -# Limit results to customers with no orders -df = df.join(df_orders, left_on="c_custkey", right_on="o_custkey", how="anti") - -# Count up the customers and the balances -df = df.aggregate( - [col("cntrycode")], - [ - F.count(col("c_custkey")).alias("numcust"), - F.sum(col("c_acctbal")).alias("totacctbal"), - ], +# Country code is the two-digit prefix of the phone number. +nation_codes = [lit(str(n)) for n in NATION_CODES] + +# Start from customers with a positive balance in one of the target country +# codes, then attach the grand-mean balance via a whole-frame window so we +# can filter per row — DataFrame stand-in for the SQL's scalar ``(select +# avg(c_acctbal) ... )`` subquery. +whole_frame = WindowFrame("rows", None, None) + +df = ( + df_customer.with_column("cntrycode", F.left(col("c_phone"), lit(2))) + .filter( + col("c_acctbal") > 0.0, + F.in_list(col("cntrycode"), nation_codes), + ) + .with_column( + "avg_balance", + F.avg(col("c_acctbal")).over(Window(window_frame=whole_frame)), + ) + .filter(col("c_acctbal") > col("avg_balance")) + # Keep only customers with no orders (anti join = NOT EXISTS). + .join(df_orders, left_on="c_custkey", right_on="o_custkey", how="anti") + .aggregate( + ["cntrycode"], + [ + F.count_star().alias("numcust"), + F.sum(col("c_acctbal")).alias("totacctbal"), + ], + ) + .sort_by("cntrycode") ) -df = df.sort(col("cntrycode").sort()) - df.show() From e0284c6e788b6fc893495ed929b9badef1cf925c Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 24 Apr 2026 13:09:24 -0400 Subject: [PATCH 29/83] feat: add AI skill to find and improve the Pythonic interface to functions (#1484) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: accept native Python types in function arguments instead of requiring lit() Update 47 functions in functions.py to accept native Python types (int, float, str) for arguments that are contextually literals, eliminating verbose lit() wrapping. For example, users can now write split_part(col("a"), ",", 2) instead of split_part(col("a"), lit(","), lit(2)). All changes are backward compatible. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: update alias function signatures to match pythonic primary functions Update instr and position (aliases of strpos) to accept Expr | str for the substring parameter, matching the updated primary function signature. Co-Authored-By: Claude Opus 4.6 (1M context) * docs: update make-pythonic skill to require alias type hint updates Alias functions that delegate to a primary function must have their type hints updated to match, even though coercion logic is only added to the primary. Added a new Step 3 to the implementation workflow for this. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: address review feedback on pythonic skill and function signatures Update SKILL.md to prevent three classes of issues: clarify that float already accepts int per PEP 484 (avoiding redundant int | float that fails ruff PYI041), add backward-compat rule for Category B so existing Expr params aren't removed, and add guidance for inline coercion with many optional nullable params instead of local helpers. Replace regexp_instr's _to_raw() helper with inline coercion matching the pattern used throughout the rest of the file. Co-Authored-By: Claude Opus 4.6 (1M context) * refactor: add coerce_to_expr helpers and replace inline coercion patterns Introduce coerce_to_expr() and coerce_to_expr_or_none() in expr.py as the complement to ensure_expr() — where ensure_expr rejects non-Expr values, these helpers wrap them via Expr.literal(). Replaces ~60 inline isinstance checks in functions.py with single-line helper calls, and updates the make-pythonic skill to document the new pattern. Co-Authored-By: Claude Opus 4.6 (1M context) * docs: add aggregate function literal detection to make-pythonic skill Add Technique 1a to detect literal-only arguments in aggregate functions. Unlike scalar UDFs which enforce literals in invoke_with_args(), aggregate functions enforce them in accumulator() via get_scalar_value(), validate_percentile_expr(), or downcast_ref::(). Without this technique, the skill would incorrectly classify arguments like approx_percentile_cont's percentile as Category A (Expr | float) when they should be Category B (float only). Updates the decision flow to branch on scalar vs aggregate before checking for literal enforcement. Co-Authored-By: Claude Opus 4.6 (1M context) * docs: add window function literal detection to make-pythonic skill Add Technique 1b to detect literal-only arguments in window functions. Window functions enforce literals in partition_evaluator() via get_scalar_value_from_args() / downcast_ref::(), not in invoke_with_args() (scalar) or accumulator() (aggregate). Updates the decision flow to branch on scalar vs aggregate vs window. Known window functions with literal-only arguments: ntile (n), lead/lag (offset, default_value), nth_value (n). Co-Authored-By: Claude Opus 4.6 (1M context) * fix: use explicit None checks, widen numeric type hints, and add tests Replace 7 fragile truthiness checks (x.expr if x else None) with explicit is not None checks to prevent silent None when zero-valued literals are passed. Widen log/power/pow type hints to Expr | int | float with noqa: PYI041 for clarity. Add unit tests for coerce_to_expr helpers and integration tests for pythonic calling conventions. Co-Authored-By: Claude Opus 4.6 (1M context) * chore: suppress FBT003 in tests and remove redundant noqa comments Add FBT003 (boolean positional value) to the per-file-ignores for python/tests/* in pyproject.toml, and remove the 6 now-redundant inline noqa: FBT003 comments across test_expr.py and test_context.py. Co-Authored-By: Claude Opus 4.6 (1M context) * docs: replace static function lists with discovery instructions in skill Replace hardcoded "Known aggregate/window functions with literal-only arguments" lists with instructions to discover them dynamically by searching the upstream crate source. Keeps a few examples as validation anchors so the agent knows its search is working correctly. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: make interrupt test reliable on Python 3.11 PyThreadState_SetAsyncExc only delivers exceptions when the thread is executing Python bytecode, not while in native (Rust/C) code. The previous test had two issues causing flakiness on Python 3.11: 1. The interrupt fired before df.collect() entered the UDF, while the thread was still in native code where async exceptions are ignored. 2. time.sleep(2.0) is a single C call where async exceptions are not checked — they're only checked between bytecode instructions. Fix by adding a threading.Event so the interrupt waits until the UDF is actually executing Python code, and by sleeping in small increments so the eval loop has opportunities to check for pending exceptions. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- .ai/skills/make-pythonic/SKILL.md | 430 +++++++++++++++++++++++++++++ pyproject.toml | 1 + python/datafusion/expr.py | 41 +++ python/datafusion/functions.py | 445 +++++++++++++++++------------- python/tests/test_context.py | 6 +- python/tests/test_dataframe.py | 21 +- python/tests/test_expr.py | 49 +++- python/tests/test_functions.py | 93 +++++++ 8 files changed, 893 insertions(+), 193 deletions(-) create mode 100644 .ai/skills/make-pythonic/SKILL.md diff --git a/.ai/skills/make-pythonic/SKILL.md b/.ai/skills/make-pythonic/SKILL.md new file mode 100644 index 000000000..57145ac6c --- /dev/null +++ b/.ai/skills/make-pythonic/SKILL.md @@ -0,0 +1,430 @@ + + +--- +name: make-pythonic +description: Audit and improve datafusion-python functions to accept native Python types (int, float, str, bool) instead of requiring explicit lit() or col() wrapping. Analyzes function signatures, checks upstream Rust implementations for type constraints, and applies the appropriate coercion pattern. +argument-hint: [scope] (e.g., "string functions", "datetime functions", "array functions", "math functions", "all", or a specific function name like "split_part") +--- + +# Make Python API Functions More Pythonic + +You are improving the datafusion-python API to feel more natural to Python users. The goal is to allow functions to accept native Python types (int, float, str, bool, etc.) for arguments that are contextually always or typically literal values, instead of requiring users to manually wrap them in `lit()`. + +**Core principle:** A Python user should be able to write `split_part(col("a"), ",", 2)` instead of `split_part(col("a"), lit(","), lit(2))` when the arguments are contextually obvious literals. + +## How to Identify Candidates + +The user may specify a scope via `$ARGUMENTS`. If no scope is given or "all" is specified, audit all functions in `python/datafusion/functions.py`. + +For each function, determine if any parameter can accept native Python types by evaluating **two complementary signals**: + +### Signal 1: Contextual Understanding + +Some arguments are contextually always or almost always literal values based on what the function does: + +| Context | Typical Arguments | Examples | +|---------|------------------|----------| +| **String position/count** | Character counts, indices, repetition counts | `left(str, n)`, `right(str, n)`, `repeat(str, n)`, `lpad(str, count, ...)` | +| **Delimiters/separators** | Fixed separator characters | `split_part(str, delim, idx)`, `concat_ws(sep, ...)` | +| **Search/replace patterns** | Literal search strings, replacements | `replace(str, from, to)`, `regexp_replace(str, pattern, replacement, flags)` | +| **Date/time parts** | Part names from a fixed set | `date_part(part, date)`, `date_trunc(part, date)` | +| **Rounding precision** | Decimal place counts | `round(val, places)`, `trunc(val, places)` | +| **Fill characters** | Padding characters | `lpad(str, count, fill)`, `rpad(str, count, fill)` | + +### Signal 2: Upstream Rust Implementation + +Check the Rust binding in `crates/core/src/functions.rs` and the upstream DataFusion function implementation to determine type constraints. The upstream source is cached locally at: + +``` +~/.cargo/registry/src/index.crates.io-*/datafusion-functions-/src/ +``` + +Check the DataFusion version in `crates/core/Cargo.toml` to find the right directory. Key subdirectories: `string/`, `datetime/`, `math/`, `regex/`. + +For **aggregate functions**, the upstream source is in a separate crate: + +``` +~/.cargo/registry/src/index.crates.io-*/datafusion-functions-aggregate-/src/ +``` + +There are five concrete techniques to check, in order of signal strength: + +#### Technique 1: Check `invoke_with_args()` for literal-only enforcement (strongest signal) + +Some functions pattern-match on `ColumnarValue::Scalar` in their `invoke_with_args()` method and **return an error** if the argument is a column/array. This means the argument **must** be a literal — passing a column expression will fail at runtime. + +Example from `date_trunc.rs`: +```rust +let granularity_str = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) = granularity { + v.to_lowercase() +} else { + return exec_err!("Granularity of `date_trunc` must be non-null scalar Utf8"); +}; +``` + +**If you find this pattern:** The argument is **Category B** — accept only the corresponding native Python type (e.g., `str`), not `Expr`. The function will error at runtime with a column expression anyway. + +#### Technique 1a: Check `accumulator()` for literal-only enforcement (aggregate functions) + +Technique 1 applies to scalar UDFs. Aggregate functions do not have `invoke_with_args()` — instead, they enforce literal-only arguments in their `accumulator()` (or `create_accumulator()`) method, which runs at planning time before any data is processed. + +Look for these patterns inside `accumulator()`: + +- `get_scalar_value(expr)` — evaluates the expression against an empty batch and errors if it's not a scalar +- `validate_percentile_expr(expr)` — specific helper used by percentile functions +- `downcast_ref::()` — checks that the physical expression is a literal constant + +Example from `approx_percentile_cont.rs`: +```rust +fn accumulator(&self, args: AccumulatorArgs) -> Result { + let percentile = + validate_percentile_expr(&args.exprs[1], "APPROX_PERCENTILE_CONT")?; + // ... +} +``` + +Where `validate_percentile_expr` calls `get_scalar_value` and errors with `"must be a literal"`. + +Example from `string_agg.rs`: +```rust +fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { + let Some(lit) = acc_args.exprs[1].as_any().downcast_ref::() else { + return not_impl_err!( + "The second argument of the string_agg function must be a string literal" + ); + }; + // ... +} +``` + +**If you find this pattern:** The argument is **Category B** — accept only the corresponding native Python type, not `Expr`. The function will error at planning time with a non-literal expression. + +To discover which aggregate functions have literal-only arguments, search the upstream aggregate crate for `get_scalar_value`, `validate_percentile_expr`, and `downcast_ref::()` inside `accumulator()` methods. For example, you should expect to find `approx_percentile_cont` (percentile) and `string_agg` (delimiter) among the results. + +#### Technique 1b: Check `partition_evaluator()` for literal-only enforcement (window functions) + +Window functions do not have `invoke_with_args()` or `accumulator()`. Instead, they enforce literal-only arguments in their `partition_evaluator()` method, which constructs the evaluator that processes each partition. + +The upstream source is in a separate crate: + +``` +~/.cargo/registry/src/index.crates.io-*/datafusion-functions-window-/src/ +``` + +Look for `get_scalar_value_from_args()` calls inside `partition_evaluator()`. This helper (defined in the window crate's `utils.rs`) calls `downcast_ref::()` and errors with `"There is only support Literal types for field at idx: {index} in Window Function"`. + +Example from `ntile.rs`: +```rust +fn partition_evaluator( + &self, + partition_evaluator_args: PartitionEvaluatorArgs, +) -> Result> { + let scalar_n = + get_scalar_value_from_args(partition_evaluator_args.input_exprs(), 0)? + .ok_or_else(|| { + exec_datafusion_err!("NTILE requires a positive integer") + })?; + // ... +} +``` + +**If you find this pattern:** The argument is **Category B** — accept only the corresponding native Python type, not `Expr`. The function will error at planning time with a non-literal expression. + +To discover which window functions have literal-only arguments, search the upstream window crate for `get_scalar_value_from_args` inside `partition_evaluator()` methods. For example, you should expect to find `ntile` (n) and `lead`/`lag` (offset, default_value) among the results. + +#### Technique 2: Check the `Signature` for data type constraints + +Each function defines a `Signature::coercible(...)` that specifies what data types each argument accepts, using `Coercion` entries. This tells you the expected **data type** even if it doesn't enforce literal-only. + +Example from `repeat.rs`: +```rust +signature: Signature::coercible( + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ), + ], + Volatility::Immutable, +), +``` + +This tells you arg 2 (`n`) must be an integer type coerced to Int64. Use this to choose the correct Python type (e.g., `int` not `str` or `float`). + +Common mappings: +| Rust Type Constraint | Python Type | +|---------------------|-------------| +| `logical_int64()` / `TypeSignatureClass::Integer` | `int` | +| `logical_float64()` / `TypeSignatureClass::Numeric` | `int \| float` | +| `logical_string()` / `TypeSignatureClass::String` | `str` | +| `LogicalType::Boolean` | `bool` | + +**Important:** In Python's type system (PEP 484), `float` already accepts `int` values, so `int | float` is redundant and will fail the `ruff` linter (rule PYI041). Use `float` alone when the Rust side accepts a float/numeric type — Python users can still pass integer literals like `log(10, col("a"))` or `power(col("a"), 3)` without issue. Only use `int` when the Rust side strictly requires an integer (e.g., `logical_int64()`). + +#### Technique 3: Check `return_field_from_args()` for `scalar_arguments` usage + +Functions that inspect literal values at query planning time use `args.scalar_arguments.get(n)` in their `return_field_from_args()` method. This indicates the argument is **expected to be a literal** for optimal behavior (e.g., to determine output type precision), but may still work as a column. + +Example from `round.rs`: +```rust +let decimal_places: Option = match args.scalar_arguments.get(1) { + None => Some(0), + Some(None) => None, // argument is not a literal (column) + Some(Some(scalar)) if scalar.is_null() => Some(0), + Some(Some(scalar)) => Some(decimal_places_from_scalar(scalar)?), +}; +``` + +**If you find this pattern:** The argument is **Category A** — accept native types AND `Expr`. It works as a column but is primarily used as a literal. + +#### Decision flow + +``` +What kind of function is this? + Scalar UDF: + Is argument rejected at runtime if not a literal? + (check invoke_with_args for ColumnarValue::Scalar-only match + exec_err!) + → YES: Category B — accept only native type, no Expr + → NO: continue below + Aggregate: + Is argument rejected at planning time if not a literal? + (check accumulator() for get_scalar_value / validate_percentile_expr / + downcast_ref::() + error) + → YES: Category B — accept only native type, no Expr + → NO: continue below + Window: + Is argument rejected at planning time if not a literal? + (check partition_evaluator() for get_scalar_value_from_args / + downcast_ref::() + error) + → YES: Category B — accept only native type, no Expr + → NO: continue below + +Does the Signature constrain it to a specific data type? + → YES: Category A — accept Expr | + → NO: Leave as Expr only +``` + +## Coercion Categories + +When making a function more pythonic, apply the correct coercion pattern based on **what the argument represents**: + +### Category A: Arguments That Should Accept Native Types AND Expr + +These are arguments that are *typically* literals but *could* be column references in advanced use cases. For these, accept a union type and coerce native types to `Expr.literal()`. + +**Type hint pattern:** `Expr | int`, `Expr | str`, `Expr | int | str`, etc. + +**When to use:** When the argument could plausibly come from a column in some use case (e.g., the repeat count might come from a column in a data-driven scenario). + +```python +def repeat(string: Expr, n: Expr | int) -> Expr: + """Repeats the ``string`` to ``n`` times. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["ha"]}) + >>> result = df.select( + ... dfn.functions.repeat(dfn.col("a"), 3).alias("r")) + >>> result.collect_column("r")[0].as_py() + 'hahaha' + """ + if not isinstance(n, Expr): + n = Expr.literal(n) + return Expr(f.repeat(string.expr, n.expr)) +``` + +### Category B: Arguments That Should ONLY Accept Specific Native Types + +These are arguments where an `Expr` never makes sense because the value must be a fixed literal known at query-planning time (not a per-row value). For these, accept only the native type(s) and wrap internally. + +**Type hint pattern:** `str`, `int`, `list[str]`, etc. (no `Expr` in the union) + +**When to use:** When the argument is from a fixed enumeration or is always a compile-time constant, **AND** the parameter was not previously typed as `Expr`: +- Separator in `concat_ws` (already typed as `str` in the Rust binding) +- Index in `array_position` (already typed as `int` in the Rust binding) +- Values that the Rust implementation already accepts as native types + +**Backward compatibility rule:** If a parameter was previously typed as `Expr`, you **must** keep `Expr` in the union even if the Rust side requires a literal. Removing `Expr` would break existing user code like `date_part(lit("year"), col("a"))`. Use **Category A** instead — accept `Expr | str` — and let users who pass column expressions discover the runtime error from the Rust side. Never silently break backward compatibility. + +```python +def concat_ws(separator: str, *args: Expr) -> Expr: + """Concatenates the list ``args`` with the separator. + + ``separator`` is already typed as ``str`` in the Rust binding, so + there is no backward-compatibility concern. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"], "b": ["world"]}) + >>> result = df.select( + ... dfn.functions.concat_ws("-", dfn.col("a"), dfn.col("b")).alias("c")) + >>> result.collect_column("c")[0].as_py() + 'hello-world' + """ + args = [arg.expr for arg in args] + return Expr(f.concat_ws(separator, args)) +``` + +### Category C: Arguments That Should Accept str as Column Name + +In some contexts a string argument naturally refers to a column name rather than a literal. This is the pattern used by DataFrame methods. + +**Type hint pattern:** `Expr | str` + +**When to use:** Only when the string contextually means a column name (rare in `functions.py`, more common in DataFrame methods). + +```python +# Use _to_raw_expr() from expr.py for this pattern +from datafusion.expr import _to_raw_expr + +def some_function(column: Expr | str) -> Expr: + raw = _to_raw_expr(column) # str -> col(str) + return Expr(f.some_function(raw)) +``` + +**IMPORTANT:** In `functions.py`, string arguments almost never mean column names. Functions operate on expressions, and column references should use `col()`. Category C applies mainly to DataFrame methods and context APIs, not to scalar/aggregate/window functions. Do NOT convert string arguments to column expressions in `functions.py` unless there is a very clear reason to do so. + +## Implementation Steps + +For each function being updated: + +### Step 1: Analyze the Function + +1. Read the current Python function signature in `python/datafusion/functions.py` +2. Read the Rust binding in `crates/core/src/functions.rs` +3. Optionally check the upstream DataFusion docs for the function +4. Determine which category (A, B, or C) applies to each parameter + +### Step 2: Update the Python Function + +1. **Change the type hints** to accept native types (e.g., `Expr` -> `Expr | int`) +2. **Add coercion logic** at the top of the function body +3. **Update the docstring** examples to use the simpler calling convention +4. **Preserve backward compatibility** — existing code using `Expr` must still work + +### Step 3: Update Alias Type Hints + +After updating a primary function, find all alias functions that delegate to it (e.g., `instr` and `position` delegate to `strpos`). Update each alias's **parameter type hints** to match the primary function's new signature. Do not add coercion logic to aliases — the primary function handles that. + +### Step 4: Update Docstring Examples (primary functions only) + +Per the project's CLAUDE.md rules: +- Every function must have doctest-style examples +- Optional parameters need examples both without and with the optional args, using keyword argument syntax +- Reuse the same input data across examples where possible + +**Update examples to demonstrate the pythonic calling convention:** + +```python +# BEFORE (old style - still works but verbose) +dfn.functions.left(dfn.col("a"), dfn.lit(3)) + +# AFTER (new style - shown in examples) +dfn.functions.left(dfn.col("a"), 3) +``` + +### Step 5: Run Tests + +After making changes, run the doctests to verify: +```bash +python -m pytest --doctest-modules python/datafusion/functions.py -v +``` + +## Coercion Helper Pattern + +Use the coercion helpers from `datafusion.expr` to convert native Python values to `Expr`. These are the complement of `ensure_expr()` — where `ensure_expr` *rejects* non-`Expr` values, the coercion helpers *wrap* them via `Expr.literal()`. + +**For required parameters** use `coerce_to_expr`: + +```python +from datafusion.expr import coerce_to_expr + +def left(string: Expr, n: Expr | int) -> Expr: + n = coerce_to_expr(n) + return Expr(f.left(string.expr, n.expr)) +``` + +**For optional nullable parameters** use `coerce_to_expr_or_none`: + +```python +from datafusion.expr import coerce_to_expr, coerce_to_expr_or_none + +def regexp_count( + string: Expr, + pattern: Expr | str, + start: Expr | int | None = None, + flags: Expr | str | None = None, +) -> Expr: + pattern = coerce_to_expr(pattern) + start = coerce_to_expr_or_none(start) + flags = coerce_to_expr_or_none(flags) + return Expr( + f.regexp_count( + string.expr, + pattern.expr, + start.expr if start is not None else None, + flags.expr if flags is not None else None, + ) + ) +``` + +Both helpers are defined in `python/datafusion/expr.py` alongside `ensure_expr`. Import them in `functions.py` via: + +```python +from datafusion.expr import coerce_to_expr, coerce_to_expr_or_none +``` + +## What NOT to Change + +- **Do not change arguments that represent data columns.** If an argument is the primary data being operated on (e.g., the `string` in `left(string, n)` or the `array` in `array_sort(array)`), it should remain `Expr` only. Users should use `col()` for column references. +- **Do not change variadic `*args: Expr` parameters.** These represent multiple expressions and should stay as `Expr`. +- **Do not change arguments where the coercion is ambiguous.** If it is unclear whether a string should be a column name or a literal, leave it as `Expr` and let the user be explicit. +- **Do not add coercion logic to simple aliases.** If a function is just `return other_function(...)`, the primary function handles coercion. However, you **must update the alias's type hints** to match the primary function's signature so that type checkers and documentation accurately reflect what the alias accepts. +- **Do not change the Rust bindings.** All coercion happens in the Python layer. The Rust functions continue to accept `PyExpr`. + +## Priority Order + +When auditing functions, process them in this order: + +1. **Date/time functions** — `date_part`, `date_trunc`, `date_bin` — these have the clearest literal arguments +2. **String functions** — `left`, `right`, `repeat`, `lpad`, `rpad`, `split_part`, `substring`, `replace`, `regexp_replace`, `regexp_match`, `regexp_count` — common and verbose without coercion +3. **Math functions** — `round`, `trunc`, `power` — numeric literal arguments +4. **Array functions** — `array_slice`, `array_position`, `array_remove_n`, `array_replace_n`, `array_resize`, `array_element` — index and count arguments +5. **Other functions** — any remaining functions with literal arguments + +## Output Format + +For each function analyzed, report: + +``` +## [Function Name] + +**Current signature:** `function(arg1: Expr, arg2: Expr) -> Expr` +**Proposed signature:** `function(arg1: Expr, arg2: Expr | int) -> Expr` +**Category:** A (accepts native + Expr) +**Arguments changed:** +- `arg2`: Expr -> Expr | int (always a literal count) +**Rust binding:** Takes PyExpr, wraps to literal internally +**Status:** [Changed / Skipped / Needs Discussion] +``` + +If asked to implement (not just audit), make the changes directly and show a summary of what was updated. diff --git a/pyproject.toml b/pyproject.toml index 327199d1a..951f7adc3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -111,6 +111,7 @@ extend-allowed-calls = ["datafusion.lit", "lit"] "ARG", "BLE001", "D", + "FBT003", "PD", "PLC0415", "PLR0913", diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 1ff6976f7..0f7f3ab5a 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -243,6 +243,8 @@ "WindowExpr", "WindowFrame", "WindowFrameBound", + "coerce_to_expr", + "coerce_to_expr_or_none", "ensure_expr", "ensure_expr_list", ] @@ -255,6 +257,10 @@ def ensure_expr(value: Expr | Any) -> expr_internal.Expr: higher level APIs consistently require explicit :func:`~datafusion.col` or :func:`~datafusion.lit` expressions. + See Also: + :func:`coerce_to_expr` — the opposite behavior: *wraps* non-``Expr`` + values as literals instead of rejecting them. + Args: value: Candidate expression or other object. @@ -299,6 +305,41 @@ def _iter( return list(_iter(exprs)) +def coerce_to_expr(value: Any) -> Expr: + """Coerce a native Python value to an ``Expr`` literal, passing ``Expr`` through. + + This is the complement of :func:`ensure_expr`: where ``ensure_expr`` + *rejects* non-``Expr`` values, ``coerce_to_expr`` *wraps* them via + :meth:`Expr.literal` so that functions can accept native Python types + (``int``, ``float``, ``str``, ``bool``, etc.) alongside ``Expr``. + + Args: + value: An ``Expr`` instance (returned as-is) or a Python literal to wrap. + + Returns: + An ``Expr`` representing the value. + """ + if isinstance(value, Expr): + return value + return Expr.literal(value) + + +def coerce_to_expr_or_none(value: Any | None) -> Expr | None: + """Coerce a value to ``Expr`` or pass ``None`` through unchanged. + + Same as :func:`coerce_to_expr` but accepts ``None`` for optional parameters. + + Args: + value: An ``Expr`` instance, a Python literal to wrap, or ``None``. + + Returns: + An ``Expr`` representing the value, or ``None``. + """ + if value is None: + return None + return coerce_to_expr(value) + + def _to_raw_expr(value: Expr | str) -> expr_internal.Expr: """Convert a Python expression or column name to its raw variant. diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 280a6d3ac..08062851a 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -49,6 +49,8 @@ Expr, SortExpr, SortKey, + coerce_to_expr, + coerce_to_expr_or_none, expr_list_to_raw_expr_list, sort_list_to_raw_sort_list, sort_or_default, @@ -383,49 +385,52 @@ def nullif(expr1: Expr, expr2: Expr) -> Expr: return Expr(f.nullif(expr1.expr, expr2.expr)) -def encode(expr: Expr, encoding: Expr) -> Expr: +def encode(expr: Expr, encoding: Expr | str) -> Expr: """Encode the ``input``, using the ``encoding``. encoding can be base64 or hex. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello"]}) >>> result = df.select( - ... dfn.functions.encode(dfn.col("a"), dfn.lit("base64")).alias("enc")) + ... dfn.functions.encode(dfn.col("a"), "base64").alias("enc")) >>> result.collect_column("enc")[0].as_py() 'aGVsbG8' """ + encoding = coerce_to_expr(encoding) return Expr(f.encode(expr.expr, encoding.expr)) -def decode(expr: Expr, encoding: Expr) -> Expr: +def decode(expr: Expr, encoding: Expr | str) -> Expr: """Decode the ``input``, using the ``encoding``. encoding can be base64 or hex. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["aGVsbG8="]}) >>> result = df.select( - ... dfn.functions.decode(dfn.col("a"), dfn.lit("base64")).alias("dec")) + ... dfn.functions.decode(dfn.col("a"), "base64").alias("dec")) >>> result.collect_column("dec")[0].as_py() b'hello' """ + encoding = coerce_to_expr(encoding) return Expr(f.decode(expr.expr, encoding.expr)) -def array_to_string(expr: Expr, delimiter: Expr) -> Expr: +def array_to_string(expr: Expr, delimiter: Expr | str) -> Expr: """Converts each element to its text representation. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) >>> result = df.select( - ... dfn.functions.array_to_string(dfn.col("a"), dfn.lit(",")).alias("s")) + ... dfn.functions.array_to_string(dfn.col("a"), ",").alias("s")) >>> result.collect_column("s")[0].as_py() '1,2,3' """ + delimiter = coerce_to_expr(delimiter) return Expr(f.array_to_string(expr.expr, delimiter.expr.cast(pa.string()))) -def array_join(expr: Expr, delimiter: Expr) -> Expr: +def array_join(expr: Expr, delimiter: Expr | str) -> Expr: """Converts each element to its text representation. See Also: @@ -434,7 +439,7 @@ def array_join(expr: Expr, delimiter: Expr) -> Expr: return array_to_string(expr, delimiter) -def list_to_string(expr: Expr, delimiter: Expr) -> Expr: +def list_to_string(expr: Expr, delimiter: Expr | str) -> Expr: """Converts each element to its text representation. See Also: @@ -443,7 +448,7 @@ def list_to_string(expr: Expr, delimiter: Expr) -> Expr: return array_to_string(expr, delimiter) -def list_join(expr: Expr, delimiter: Expr) -> Expr: +def list_join(expr: Expr, delimiter: Expr | str) -> Expr: """Converts each element to its text representation. See Also: @@ -479,7 +484,7 @@ def in_list(arg: Expr, values: list[Expr], negated: bool = False) -> Expr: return Expr(f.in_list(arg.expr, values, negated)) -def digest(value: Expr, method: Expr) -> Expr: +def digest(value: Expr, method: Expr | str) -> Expr: """Computes the binary hash of an expression using the specified algorithm. Standard algorithms are md5, sha224, sha256, sha384, sha512, blake2s, @@ -489,24 +494,26 @@ def digest(value: Expr, method: Expr) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello"]}) >>> result = df.select( - ... dfn.functions.digest(dfn.col("a"), dfn.lit("md5")).alias("d")) + ... dfn.functions.digest(dfn.col("a"), "md5").alias("d")) >>> len(result.collect_column("d")[0].as_py()) > 0 True """ + method = coerce_to_expr(method) return Expr(f.digest(value.expr, method.expr)) -def contains(string: Expr, search_str: Expr) -> Expr: +def contains(string: Expr, search_str: Expr | str) -> Expr: """Returns true if ``search_str`` is found within ``string`` (case-sensitive). Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["the quick brown fox"]}) >>> result = df.select( - ... dfn.functions.contains(dfn.col("a"), dfn.lit("brown")).alias("c")) + ... dfn.functions.contains(dfn.col("a"), "brown").alias("c")) >>> result.collect_column("c")[0].as_py() True """ + search_str = coerce_to_expr(search_str) return Expr(f.contains(string.expr, search_str.expr)) @@ -969,17 +976,18 @@ def degrees(arg: Expr) -> Expr: return Expr(f.degrees(arg.expr)) -def ends_with(arg: Expr, suffix: Expr) -> Expr: +def ends_with(arg: Expr, suffix: Expr | str) -> Expr: """Returns true if the ``string`` ends with the ``suffix``, false otherwise. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["abc","b","c"]}) >>> ends_with_df = df.select( - ... dfn.functions.ends_with(dfn.col("a"), dfn.lit("c")).alias("ends_with")) + ... dfn.functions.ends_with(dfn.col("a"), "c").alias("ends_with")) >>> ends_with_df.collect_column("ends_with")[0].as_py() True """ + suffix = coerce_to_expr(suffix) return Expr(f.ends_with(arg.expr, suffix.expr)) @@ -1011,7 +1019,7 @@ def factorial(arg: Expr) -> Expr: return Expr(f.factorial(arg.expr)) -def find_in_set(string: Expr, string_list: Expr) -> Expr: +def find_in_set(string: Expr, string_list: Expr | str) -> Expr: """Find a string in a list of strings. Returns a value in the range of 1 to N if the string is in the string list @@ -1023,10 +1031,11 @@ def find_in_set(string: Expr, string_list: Expr) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["b"]}) >>> result = df.select( - ... dfn.functions.find_in_set(dfn.col("a"), dfn.lit("a,b,c")).alias("pos")) + ... dfn.functions.find_in_set(dfn.col("a"), "a,b,c").alias("pos")) >>> result.collect_column("pos")[0].as_py() 2 """ + string_list = coerce_to_expr(string_list) return Expr(f.find_in_set(string.expr, string_list.expr)) @@ -1102,7 +1111,7 @@ def initcap(string: Expr) -> Expr: return Expr(f.initcap(string.expr)) -def instr(string: Expr, substring: Expr) -> Expr: +def instr(string: Expr, substring: Expr | str) -> Expr: """Finds the position from where the ``substring`` matches the ``string``. See Also: @@ -1158,31 +1167,33 @@ def least(*args: Expr) -> Expr: return Expr(f.least(*exprs)) -def left(string: Expr, n: Expr) -> Expr: +def left(string: Expr, n: Expr | int) -> Expr: """Returns the first ``n`` characters in the ``string``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["the cat"]}) >>> left_df = df.select( - ... dfn.functions.left(dfn.col("a"), dfn.lit(3)).alias("left")) + ... dfn.functions.left(dfn.col("a"), 3).alias("left")) >>> left_df.collect_column("left")[0].as_py() 'the' """ + n = coerce_to_expr(n) return Expr(f.left(string.expr, n.expr)) -def levenshtein(string1: Expr, string2: Expr) -> Expr: +def levenshtein(string1: Expr, string2: Expr | str) -> Expr: """Returns the Levenshtein distance between the two given strings. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["kitten"]}) >>> result = df.select( - ... dfn.functions.levenshtein(dfn.col("a"), dfn.lit("sitting")).alias("d")) + ... dfn.functions.levenshtein(dfn.col("a"), "sitting").alias("d")) >>> result.collect_column("d")[0].as_py() 3 """ + string2 = coerce_to_expr(string2) return Expr(f.levenshtein(string1.expr, string2.expr)) @@ -1199,18 +1210,19 @@ def ln(arg: Expr) -> Expr: return Expr(f.ln(arg.expr)) -def log(base: Expr, num: Expr) -> Expr: +def log(base: Expr | int | float, num: Expr) -> Expr: # noqa: PYI041 """Returns the logarithm of a number for a particular ``base``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [100.0]}) >>> result = df.select( - ... dfn.functions.log(dfn.lit(10.0), dfn.col("a")).alias("log") + ... dfn.functions.log(10.0, dfn.col("a")).alias("log") ... ) >>> result.collect_column("log")[0].as_py() 2.0 """ + base = coerce_to_expr(base) return Expr(f.log(base.expr, num.expr)) @@ -1253,7 +1265,7 @@ def lower(arg: Expr) -> Expr: return Expr(f.lower(arg.expr)) -def lpad(string: Expr, count: Expr, characters: Expr | None = None) -> Expr: +def lpad(string: Expr, count: Expr | int, characters: Expr | str | None = None) -> Expr: """Add left padding to a string. Extends the string to length length by prepending the characters fill (a @@ -1264,9 +1276,7 @@ def lpad(string: Expr, count: Expr, characters: Expr | None = None) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["the cat", "a hat"]}) >>> lpad_df = df.select( - ... dfn.functions.lpad( - ... dfn.col("a"), dfn.lit(6) - ... ).alias("lpad")) + ... dfn.functions.lpad(dfn.col("a"), 6).alias("lpad")) >>> lpad_df.collect_column("lpad")[0].as_py() 'the ca' >>> lpad_df.collect_column("lpad")[1].as_py() @@ -1274,12 +1284,13 @@ def lpad(string: Expr, count: Expr, characters: Expr | None = None) -> Expr: >>> result = df.select( ... dfn.functions.lpad( - ... dfn.col("a"), dfn.lit(10), characters=dfn.lit(".") + ... dfn.col("a"), 10, characters="." ... ).alias("lpad")) >>> result.collect_column("lpad")[0].as_py() '...the cat' """ - characters = characters if characters is not None else Expr.literal(" ") + count = coerce_to_expr(count) + characters = coerce_to_expr(characters if characters is not None else " ") return Expr(f.lpad(string.expr, count.expr, characters.expr)) @@ -1374,7 +1385,10 @@ def octet_length(arg: Expr) -> Expr: def overlay( - string: Expr, substring: Expr, start: Expr, length: Expr | None = None + string: Expr, + substring: Expr | str, + start: Expr | int, + length: Expr | int | None = None, ) -> Expr: """Replace a substring with a new substring. @@ -1385,13 +1399,15 @@ def overlay( >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["abcdef"]}) >>> result = df.select( - ... dfn.functions.overlay(dfn.col("a"), dfn.lit("XY"), dfn.lit(3), - ... dfn.lit(2)).alias("o")) + ... dfn.functions.overlay(dfn.col("a"), "XY", 3, 2).alias("o")) >>> result.collect_column("o")[0].as_py() 'abXYef' """ + substring = coerce_to_expr(substring) + start = coerce_to_expr(start) if length is None: return Expr(f.overlay(string.expr, substring.expr, start.expr)) + length = coerce_to_expr(length) return Expr(f.overlay(string.expr, substring.expr, start.expr, length.expr)) @@ -1411,7 +1427,7 @@ def pi() -> Expr: return Expr(f.pi()) -def position(string: Expr, substring: Expr) -> Expr: +def position(string: Expr, substring: Expr | str) -> Expr: """Finds the position from where the ``substring`` matches the ``string``. See Also: @@ -1420,22 +1436,23 @@ def position(string: Expr, substring: Expr) -> Expr: return strpos(string, substring) -def power(base: Expr, exponent: Expr) -> Expr: +def power(base: Expr, exponent: Expr | int | float) -> Expr: # noqa: PYI041 """Returns ``base`` raised to the power of ``exponent``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [2.0]}) >>> result = df.select( - ... dfn.functions.power(dfn.col("a"), dfn.lit(3.0)).alias("pow") + ... dfn.functions.power(dfn.col("a"), 3.0).alias("pow") ... ) >>> result.collect_column("pow")[0].as_py() 8.0 """ + exponent = coerce_to_expr(exponent) return Expr(f.power(base.expr, exponent.expr)) -def pow(base: Expr, exponent: Expr) -> Expr: +def pow(base: Expr, exponent: Expr | int | float) -> Expr: # noqa: PYI041 """Returns ``base`` raised to the power of ``exponent``. See Also: @@ -1460,7 +1477,9 @@ def radians(arg: Expr) -> Expr: return Expr(f.radians(arg.expr)) -def regexp_like(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: +def regexp_like( + string: Expr, regex: Expr | str, flags: Expr | str | None = None +) -> Expr: r"""Find if any regular expression (regex) matches exist. Tests a string using a regular expression returning true if at least one match, @@ -1470,9 +1489,7 @@ def regexp_like(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello123"]}) >>> result = df.select( - ... dfn.functions.regexp_like( - ... dfn.col("a"), dfn.lit("\\d+") - ... ).alias("m") + ... dfn.functions.regexp_like(dfn.col("a"), "\\d+").alias("m") ... ) >>> result.collect_column("m")[0].as_py() True @@ -1481,19 +1498,24 @@ def regexp_like(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: >>> result = df.select( ... dfn.functions.regexp_like( - ... dfn.col("a"), dfn.lit("HELLO"), - ... flags=dfn.lit("i"), + ... dfn.col("a"), "HELLO", flags="i", ... ).alias("m") ... ) >>> result.collect_column("m")[0].as_py() True """ - if flags is not None: - flags = flags.expr - return Expr(f.regexp_like(string.expr, regex.expr, flags)) + regex = coerce_to_expr(regex) + flags = coerce_to_expr_or_none(flags) + return Expr( + f.regexp_like( + string.expr, regex.expr, flags.expr if flags is not None else None + ) + ) -def regexp_match(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: +def regexp_match( + string: Expr, regex: Expr | str, flags: Expr | str | None = None +) -> Expr: r"""Perform regular expression (regex) matching. Returns an array with each element containing the leftmost-first match of the @@ -1503,9 +1525,7 @@ def regexp_match(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello 42 world"]}) >>> result = df.select( - ... dfn.functions.regexp_match( - ... dfn.col("a"), dfn.lit("(\\d+)") - ... ).alias("m") + ... dfn.functions.regexp_match(dfn.col("a"), "(\\d+)").alias("m") ... ) >>> result.collect_column("m")[0].as_py() ['42'] @@ -1514,20 +1534,26 @@ def regexp_match(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: >>> result = df.select( ... dfn.functions.regexp_match( - ... dfn.col("a"), dfn.lit("(HELLO)"), - ... flags=dfn.lit("i"), + ... dfn.col("a"), "(HELLO)", flags="i", ... ).alias("m") ... ) >>> result.collect_column("m")[0].as_py() ['hello'] """ - if flags is not None: - flags = flags.expr - return Expr(f.regexp_match(string.expr, regex.expr, flags)) + regex = coerce_to_expr(regex) + flags = coerce_to_expr_or_none(flags) + return Expr( + f.regexp_match( + string.expr, regex.expr, flags.expr if flags is not None else None + ) + ) def regexp_replace( - string: Expr, pattern: Expr, replacement: Expr, flags: Expr | None = None + string: Expr, + pattern: Expr | str, + replacement: Expr | str, + flags: Expr | str | None = None, ) -> Expr: r"""Replaces substring(s) matching a PCRE-like regular expression. @@ -1542,8 +1568,7 @@ def regexp_replace( >>> df = ctx.from_pydict({"a": ["hello 42"]}) >>> result = df.select( ... dfn.functions.regexp_replace( - ... dfn.col("a"), dfn.lit("\\d+"), - ... dfn.lit("XX") + ... dfn.col("a"), "\\d+", "XX" ... ).alias("r") ... ) >>> result.collect_column("r")[0].as_py() @@ -1554,20 +1579,30 @@ def regexp_replace( >>> df = ctx.from_pydict({"a": ["a1 b2 c3"]}) >>> result = df.select( ... dfn.functions.regexp_replace( - ... dfn.col("a"), dfn.lit("\\d+"), - ... dfn.lit("X"), flags=dfn.lit("g"), + ... dfn.col("a"), "\\d+", "X", flags="g", ... ).alias("r") ... ) >>> result.collect_column("r")[0].as_py() 'aX bX cX' """ - if flags is not None: - flags = flags.expr - return Expr(f.regexp_replace(string.expr, pattern.expr, replacement.expr, flags)) + pattern = coerce_to_expr(pattern) + replacement = coerce_to_expr(replacement) + flags = coerce_to_expr_or_none(flags) + return Expr( + f.regexp_replace( + string.expr, + pattern.expr, + replacement.expr, + flags.expr if flags is not None else None, + ) + ) def regexp_count( - string: Expr, pattern: Expr, start: Expr | None = None, flags: Expr | None = None + string: Expr, + pattern: Expr | str, + start: Expr | int | None = None, + flags: Expr | str | None = None, ) -> Expr: """Returns the number of matches in a string. @@ -1578,9 +1613,7 @@ def regexp_count( >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["abcabc"]}) >>> result = df.select( - ... dfn.functions.regexp_count( - ... dfn.col("a"), dfn.lit("abc") - ... ).alias("c")) + ... dfn.functions.regexp_count(dfn.col("a"), "abc").alias("c")) >>> result.collect_column("c")[0].as_py() 2 @@ -1589,25 +1622,31 @@ def regexp_count( >>> result = df.select( ... dfn.functions.regexp_count( - ... dfn.col("a"), dfn.lit("ABC"), - ... start=dfn.lit(4), flags=dfn.lit("i"), + ... dfn.col("a"), "ABC", start=4, flags="i", ... ).alias("c")) >>> result.collect_column("c")[0].as_py() 1 """ - if flags is not None: - flags = flags.expr - start = start.expr if start is not None else start - return Expr(f.regexp_count(string.expr, pattern.expr, start, flags)) + pattern = coerce_to_expr(pattern) + start = coerce_to_expr_or_none(start) + flags = coerce_to_expr_or_none(flags) + return Expr( + f.regexp_count( + string.expr, + pattern.expr, + start.expr if start is not None else None, + flags.expr if flags is not None else None, + ) + ) def regexp_instr( values: Expr, - regex: Expr, - start: Expr | None = None, - n: Expr | None = None, - flags: Expr | None = None, - sub_expr: Expr | None = None, + regex: Expr | str, + start: Expr | int | None = None, + n: Expr | int | None = None, + flags: Expr | str | None = None, + sub_expr: Expr | int | None = None, ) -> Expr: r"""Returns the position of a regular expression match in a string. @@ -1623,9 +1662,7 @@ def regexp_instr( >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello 42 world"]}) >>> result = df.select( - ... dfn.functions.regexp_instr( - ... dfn.col("a"), dfn.lit("\\d+") - ... ).alias("pos") + ... dfn.functions.regexp_instr(dfn.col("a"), "\\d+").alias("pos") ... ) >>> result.collect_column("pos")[0].as_py() 7 @@ -1636,9 +1673,8 @@ def regexp_instr( >>> df = ctx.from_pydict({"a": ["abc ABC abc"]}) >>> result = df.select( ... dfn.functions.regexp_instr( - ... dfn.col("a"), dfn.lit("abc"), - ... start=dfn.lit(2), n=dfn.lit(1), - ... flags=dfn.lit("i"), + ... dfn.col("a"), "abc", + ... start=2, n=1, flags="i", ... ).alias("pos") ... ) >>> result.collect_column("pos")[0].as_py() @@ -1648,56 +1684,58 @@ def regexp_instr( >>> result = df.select( ... dfn.functions.regexp_instr( - ... dfn.col("a"), dfn.lit("(abc)"), - ... sub_expr=dfn.lit(1), + ... dfn.col("a"), "(abc)", sub_expr=1, ... ).alias("pos") ... ) >>> result.collect_column("pos")[0].as_py() 1 """ - start = start.expr if start is not None else None - n = n.expr if n is not None else None - flags = flags.expr if flags is not None else None - sub_expr = sub_expr.expr if sub_expr is not None else None + regex = coerce_to_expr(regex) + start = coerce_to_expr_or_none(start) + n = coerce_to_expr_or_none(n) + flags = coerce_to_expr_or_none(flags) + sub_expr = coerce_to_expr_or_none(sub_expr) return Expr( f.regexp_instr( values.expr, regex.expr, - start, - n, - flags, - sub_expr, + start.expr if start is not None else None, + n.expr if n is not None else None, + flags.expr if flags is not None else None, + sub_expr.expr if sub_expr is not None else None, ) ) -def repeat(string: Expr, n: Expr) -> Expr: +def repeat(string: Expr, n: Expr | int) -> Expr: """Repeats the ``string`` to ``n`` times. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["ha"]}) >>> result = df.select( - ... dfn.functions.repeat(dfn.col("a"), dfn.lit(3)).alias("r")) + ... dfn.functions.repeat(dfn.col("a"), 3).alias("r")) >>> result.collect_column("r")[0].as_py() 'hahaha' """ + n = coerce_to_expr(n) return Expr(f.repeat(string.expr, n.expr)) -def replace(string: Expr, from_val: Expr, to_val: Expr) -> Expr: +def replace(string: Expr, from_val: Expr | str, to_val: Expr | str) -> Expr: """Replaces all occurrences of ``from_val`` with ``to_val`` in the ``string``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello world"]}) >>> result = df.select( - ... dfn.functions.replace(dfn.col("a"), dfn.lit("world"), - ... dfn.lit("there")).alias("r")) + ... dfn.functions.replace(dfn.col("a"), "world", "there").alias("r")) >>> result.collect_column("r")[0].as_py() 'hello there' """ + from_val = coerce_to_expr(from_val) + to_val = coerce_to_expr(to_val) return Expr(f.replace(string.expr, from_val.expr, to_val.expr)) @@ -1714,39 +1752,39 @@ def reverse(arg: Expr) -> Expr: return Expr(f.reverse(arg.expr)) -def right(string: Expr, n: Expr) -> Expr: +def right(string: Expr, n: Expr | int) -> Expr: """Returns the last ``n`` characters in the ``string``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello"]}) - >>> result = df.select(dfn.functions.right(dfn.col("a"), dfn.lit(3)).alias("r")) + >>> result = df.select(dfn.functions.right(dfn.col("a"), 3).alias("r")) >>> result.collect_column("r")[0].as_py() 'llo' """ + n = coerce_to_expr(n) return Expr(f.right(string.expr, n.expr)) -def round(value: Expr, decimal_places: Expr | None = None) -> Expr: +def round(value: Expr, decimal_places: Expr | int | None = None) -> Expr: """Round the argument to the nearest integer. If the optional ``decimal_places`` is specified, round to the nearest number of decimal places. You can specify a negative number of decimal places. For example - ``round(lit(125.2345), lit(-2))`` would yield a value of ``100.0``. + ``round(lit(125.2345), -2)`` would yield a value of ``100.0``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [1.567]}) - >>> result = df.select(dfn.functions.round(dfn.col("a"), dfn.lit(2)).alias("r")) + >>> result = df.select(dfn.functions.round(dfn.col("a"), 2).alias("r")) >>> result.collect_column("r")[0].as_py() 1.57 """ - if decimal_places is None: - decimal_places = Expr.literal(0) + decimal_places = coerce_to_expr(decimal_places if decimal_places is not None else 0) return Expr(f.round(value.expr, decimal_places.expr)) -def rpad(string: Expr, count: Expr, characters: Expr | None = None) -> Expr: +def rpad(string: Expr, count: Expr | int, characters: Expr | str | None = None) -> Expr: """Add right padding to a string. Extends the string to length length by appending the characters fill (a space @@ -1756,11 +1794,12 @@ def rpad(string: Expr, count: Expr, characters: Expr | None = None) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hi"]}) >>> result = df.select( - ... dfn.functions.rpad(dfn.col("a"), dfn.lit(5), dfn.lit("!")).alias("r")) + ... dfn.functions.rpad(dfn.col("a"), 5, "!").alias("r")) >>> result.collect_column("r")[0].as_py() 'hi!!!' """ - characters = characters if characters is not None else Expr.literal(" ") + count = coerce_to_expr(count) + characters = coerce_to_expr(characters if characters is not None else " ") return Expr(f.rpad(string.expr, count.expr, characters.expr)) @@ -1876,7 +1915,7 @@ def sinh(arg: Expr) -> Expr: return Expr(f.sinh(arg.expr)) -def split_part(string: Expr, delimiter: Expr, index: Expr) -> Expr: +def split_part(string: Expr, delimiter: Expr | str, index: Expr | int) -> Expr: """Split a string and return one part. Splits a string based on a delimiter and picks out the desired field based @@ -1886,12 +1925,12 @@ def split_part(string: Expr, delimiter: Expr, index: Expr) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["a,b,c"]}) >>> result = df.select( - ... dfn.functions.split_part( - ... dfn.col("a"), dfn.lit(","), dfn.lit(2) - ... ).alias("s")) + ... dfn.functions.split_part(dfn.col("a"), ",", 2).alias("s")) >>> result.collect_column("s")[0].as_py() 'b' """ + delimiter = coerce_to_expr(delimiter) + index = coerce_to_expr(index) return Expr(f.split_part(string.expr, delimiter.expr, index.expr)) @@ -1908,49 +1947,52 @@ def sqrt(arg: Expr) -> Expr: return Expr(f.sqrt(arg.expr)) -def starts_with(string: Expr, prefix: Expr) -> Expr: +def starts_with(string: Expr, prefix: Expr | str) -> Expr: """Returns true if string starts with prefix. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello_from_datafusion"]}) >>> result = df.select( - ... dfn.functions.starts_with(dfn.col("a"), dfn.lit("hello")).alias("sw")) + ... dfn.functions.starts_with(dfn.col("a"), "hello").alias("sw")) >>> result.collect_column("sw")[0].as_py() True """ + prefix = coerce_to_expr(prefix) return Expr(f.starts_with(string.expr, prefix.expr)) -def strpos(string: Expr, substring: Expr) -> Expr: +def strpos(string: Expr, substring: Expr | str) -> Expr: """Finds the position from where the ``substring`` matches the ``string``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello"]}) >>> result = df.select( - ... dfn.functions.strpos(dfn.col("a"), dfn.lit("llo")).alias("pos")) + ... dfn.functions.strpos(dfn.col("a"), "llo").alias("pos")) >>> result.collect_column("pos")[0].as_py() 3 """ + substring = coerce_to_expr(substring) return Expr(f.strpos(string.expr, substring.expr)) -def substr(string: Expr, position: Expr) -> Expr: +def substr(string: Expr, position: Expr | int) -> Expr: """Substring from the ``position`` to the end. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello"]}) >>> result = df.select( - ... dfn.functions.substr(dfn.col("a"), dfn.lit(3)).alias("s")) + ... dfn.functions.substr(dfn.col("a"), 3).alias("s")) >>> result.collect_column("s")[0].as_py() 'llo' """ + position = coerce_to_expr(position) return Expr(f.substr(string.expr, position.expr)) -def substr_index(string: Expr, delimiter: Expr, count: Expr) -> Expr: +def substr_index(string: Expr, delimiter: Expr | str, count: Expr | int) -> Expr: """Returns an indexed substring. The return will be the ``string`` from before ``count`` occurrences of @@ -1960,27 +2002,28 @@ def substr_index(string: Expr, delimiter: Expr, count: Expr) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["a.b.c"]}) >>> result = df.select( - ... dfn.functions.substr_index(dfn.col("a"), dfn.lit("."), - ... dfn.lit(2)).alias("s")) + ... dfn.functions.substr_index(dfn.col("a"), ".", 2).alias("s")) >>> result.collect_column("s")[0].as_py() 'a.b' """ + delimiter = coerce_to_expr(delimiter) + count = coerce_to_expr(count) return Expr(f.substr_index(string.expr, delimiter.expr, count.expr)) -def substring(string: Expr, position: Expr, length: Expr) -> Expr: +def substring(string: Expr, position: Expr | int, length: Expr | int) -> Expr: """Substring from the ``position`` with ``length`` characters. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello world"]}) >>> result = df.select( - ... dfn.functions.substring( - ... dfn.col("a"), dfn.lit(1), dfn.lit(5) - ... ).alias("s")) + ... dfn.functions.substring(dfn.col("a"), 1, 5).alias("s")) >>> result.collect_column("s")[0].as_py() 'hello' """ + position = coerce_to_expr(position) + length = coerce_to_expr(length) return Expr(f.substring(string.expr, position.expr, length.expr)) @@ -2053,7 +2096,7 @@ def current_timestamp() -> Expr: return now() -def to_char(arg: Expr, formatter: Expr) -> Expr: +def to_char(arg: Expr, formatter: Expr | str) -> Expr: """Returns a string representation of a date, time, timestamp or duration. For usage of ``formatter`` see the rust chrono package ``strftime`` package. @@ -2066,16 +2109,17 @@ def to_char(arg: Expr, formatter: Expr) -> Expr: >>> result = df.select( ... dfn.functions.to_char( ... dfn.functions.to_timestamp(dfn.col("a")), - ... dfn.lit("%Y/%m/%d"), + ... "%Y/%m/%d", ... ).alias("formatted") ... ) >>> result.collect_column("formatted")[0].as_py() '2021/01/01' """ + formatter = coerce_to_expr(formatter) return Expr(f.to_char(arg.expr, formatter.expr)) -def date_format(arg: Expr, formatter: Expr) -> Expr: +def date_format(arg: Expr, formatter: Expr | str) -> Expr: """Returns a string representation of a date, time, timestamp or duration. See Also: @@ -2287,7 +2331,7 @@ def current_time() -> Expr: return Expr(f.current_time()) -def datepart(part: Expr, date: Expr) -> Expr: +def datepart(part: Expr | str, date: Expr) -> Expr: """Return a specified part of a date. See Also: @@ -2296,22 +2340,28 @@ def datepart(part: Expr, date: Expr) -> Expr: return date_part(part, date) -def date_part(part: Expr, date: Expr) -> Expr: +def date_part(part: Expr | str, date: Expr) -> Expr: """Extracts a subfield from the date. + Args: + part: The part of the date to extract. Must be one of ``"year"``, + ``"month"``, ``"day"``, ``"hour"``, ``"minute"``, ``"second"``, etc. + date: The date expression to extract from. + Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["2021-07-15T00:00:00"]}) >>> df = df.select(dfn.functions.to_timestamp(dfn.col("a")).alias("a")) >>> result = df.select( - ... dfn.functions.date_part(dfn.lit("year"), dfn.col("a")).alias("y")) + ... dfn.functions.date_part("year", dfn.col("a")).alias("y")) >>> result.collect_column("y")[0].as_py() 2021 """ + part = coerce_to_expr(part) return Expr(f.date_part(part.expr, date.expr)) -def extract(part: Expr, date: Expr) -> Expr: +def extract(part: Expr | str, date: Expr) -> Expr: """Extracts a subfield from the date. See Also: @@ -2320,25 +2370,29 @@ def extract(part: Expr, date: Expr) -> Expr: return date_part(part, date) -def date_trunc(part: Expr, date: Expr) -> Expr: +def date_trunc(part: Expr | str, date: Expr) -> Expr: """Truncates the date to a specified level of precision. + Args: + part: The precision to truncate to. Must be one of ``"year"``, + ``"month"``, ``"day"``, ``"hour"``, ``"minute"``, ``"second"``, etc. + date: The date expression to truncate. + Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["2021-07-15T12:34:56"]}) >>> df = df.select(dfn.functions.to_timestamp(dfn.col("a")).alias("a")) >>> result = df.select( - ... dfn.functions.date_trunc( - ... dfn.lit("month"), dfn.col("a") - ... ).alias("t") + ... dfn.functions.date_trunc("month", dfn.col("a")).alias("t") ... ) >>> str(result.collect_column("t")[0].as_py()) '2021-07-01 00:00:00' """ + part = coerce_to_expr(part) return Expr(f.date_trunc(part.expr, date.expr)) -def datetrunc(part: Expr, date: Expr) -> Expr: +def datetrunc(part: Expr | str, date: Expr) -> Expr: """Truncates the date to a specified level of precision. See Also: @@ -2399,18 +2453,19 @@ def make_time(hour: Expr, minute: Expr, second: Expr) -> Expr: return Expr(f.make_time(hour.expr, minute.expr, second.expr)) -def translate(string: Expr, from_val: Expr, to_val: Expr) -> Expr: +def translate(string: Expr, from_val: Expr | str, to_val: Expr | str) -> Expr: """Replaces the characters in ``from_val`` with the counterpart in ``to_val``. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello"]}) >>> result = df.select( - ... dfn.functions.translate(dfn.col("a"), dfn.lit("helo"), - ... dfn.lit("HELO")).alias("t")) + ... dfn.functions.translate(dfn.col("a"), "helo", "HELO").alias("t")) >>> result.collect_column("t")[0].as_py() 'HELLO' """ + from_val = coerce_to_expr(from_val) + to_val = coerce_to_expr(to_val) return Expr(f.translate(string.expr, from_val.expr, to_val.expr)) @@ -2427,27 +2482,24 @@ def trim(arg: Expr) -> Expr: return Expr(f.trim(arg.expr)) -def trunc(num: Expr, precision: Expr | None = None) -> Expr: +def trunc(num: Expr, precision: Expr | int | None = None) -> Expr: """Truncate the number toward zero with optional precision. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [1.567]}) >>> result = df.select( - ... dfn.functions.trunc( - ... dfn.col("a") - ... ).alias("t")) + ... dfn.functions.trunc(dfn.col("a")).alias("t")) >>> result.collect_column("t")[0].as_py() 1.0 >>> result = df.select( - ... dfn.functions.trunc( - ... dfn.col("a"), precision=dfn.lit(2) - ... ).alias("t")) + ... dfn.functions.trunc(dfn.col("a"), precision=2).alias("t")) >>> result.collect_column("t")[0].as_py() 1.56 """ if precision is not None: + precision = coerce_to_expr(precision) return Expr(f.trunc(num.expr, precision.expr)) return Expr(f.trunc(num.expr)) @@ -2928,17 +2980,18 @@ def list_dims(array: Expr) -> Expr: return array_dims(array) -def array_element(array: Expr, n: Expr) -> Expr: +def array_element(array: Expr, n: Expr | int) -> Expr: """Extracts the element with the index n from the array. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [[10, 20, 30]]}) >>> result = df.select( - ... dfn.functions.array_element(dfn.col("a"), dfn.lit(2)).alias("result")) + ... dfn.functions.array_element(dfn.col("a"), 2).alias("result")) >>> result.collect_column("result")[0].as_py() 20 """ + n = coerce_to_expr(n) return Expr(f.array_element(array.expr, n.expr)) @@ -2964,7 +3017,7 @@ def list_empty(array: Expr) -> Expr: return array_empty(array) -def array_extract(array: Expr, n: Expr) -> Expr: +def array_extract(array: Expr, n: Expr | int) -> Expr: """Extracts the element with the index n from the array. See Also: @@ -2973,7 +3026,7 @@ def array_extract(array: Expr, n: Expr) -> Expr: return array_element(array, n) -def list_element(array: Expr, n: Expr) -> Expr: +def list_element(array: Expr, n: Expr | int) -> Expr: """Extracts the element with the index n from the array. See Also: @@ -2982,7 +3035,7 @@ def list_element(array: Expr, n: Expr) -> Expr: return array_element(array, n) -def list_extract(array: Expr, n: Expr) -> Expr: +def list_extract(array: Expr, n: Expr | int) -> Expr: """Extracts the element with the index n from the array. See Also: @@ -3332,22 +3385,24 @@ def list_remove(array: Expr, element: Expr) -> Expr: return array_remove(array, element) -def array_remove_n(array: Expr, element: Expr, max: Expr) -> Expr: +def array_remove_n(array: Expr, element: Expr, max: Expr | int) -> Expr: """Removes the first ``max`` elements from the array equal to the given value. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [[1, 2, 1, 1]]}) >>> result = df.select( - ... dfn.functions.array_remove_n(dfn.col("a"), dfn.lit(1), - ... dfn.lit(2)).alias("result")) + ... dfn.functions.array_remove_n( + ... dfn.col("a"), dfn.lit(1), 2 + ... ).alias("result")) >>> result.collect_column("result")[0].as_py() [2, 1] """ + max = coerce_to_expr(max) return Expr(f.array_remove_n(array.expr, element.expr, max.expr)) -def list_remove_n(array: Expr, element: Expr, max: Expr) -> Expr: +def list_remove_n(array: Expr, element: Expr, max: Expr | int) -> Expr: """Removes the first ``max`` elements from the array equal to the given value. See Also: @@ -3381,21 +3436,22 @@ def list_remove_all(array: Expr, element: Expr) -> Expr: return array_remove_all(array, element) -def array_repeat(element: Expr, count: Expr) -> Expr: +def array_repeat(element: Expr, count: Expr | int) -> Expr: """Returns an array containing ``element`` ``count`` times. Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [1]}) >>> result = df.select( - ... dfn.functions.array_repeat(dfn.lit(3), dfn.lit(3)).alias("result")) + ... dfn.functions.array_repeat(dfn.lit(3), 3).alias("result")) >>> result.collect_column("result")[0].as_py() [3, 3, 3] """ + count = coerce_to_expr(count) return Expr(f.array_repeat(element.expr, count.expr)) -def list_repeat(element: Expr, count: Expr) -> Expr: +def list_repeat(element: Expr, count: Expr | int) -> Expr: """Returns an array containing ``element`` ``count`` times. See Also: @@ -3428,7 +3484,7 @@ def list_replace(array: Expr, from_val: Expr, to_val: Expr) -> Expr: return array_replace(array, from_val, to_val) -def array_replace_n(array: Expr, from_val: Expr, to_val: Expr, max: Expr) -> Expr: +def array_replace_n(array: Expr, from_val: Expr, to_val: Expr, max: Expr | int) -> Expr: """Replace ``n`` occurrences of ``from_val`` with ``to_val``. Replaces the first ``max`` occurrences of the specified element with another @@ -3438,15 +3494,17 @@ def array_replace_n(array: Expr, from_val: Expr, to_val: Expr, max: Expr) -> Exp >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [[1, 2, 1, 1]]}) >>> result = df.select( - ... dfn.functions.array_replace_n(dfn.col("a"), dfn.lit(1), dfn.lit(9), - ... dfn.lit(2)).alias("result")) + ... dfn.functions.array_replace_n( + ... dfn.col("a"), dfn.lit(1), dfn.lit(9), 2 + ... ).alias("result")) >>> result.collect_column("result")[0].as_py() [9, 2, 9, 1] """ + max = coerce_to_expr(max) return Expr(f.array_replace_n(array.expr, from_val.expr, to_val.expr, max.expr)) -def list_replace_n(array: Expr, from_val: Expr, to_val: Expr, max: Expr) -> Expr: +def list_replace_n(array: Expr, from_val: Expr, to_val: Expr, max: Expr | int) -> Expr: """Replace ``n`` occurrences of ``from_val`` with ``to_val``. Replaces the first ``max`` occurrences of the specified element with another @@ -3529,7 +3587,10 @@ def list_sort(array: Expr, descending: bool = False, null_first: bool = False) - def array_slice( - array: Expr, begin: Expr, end: Expr, stride: Expr | None = None + array: Expr, + begin: Expr | int, + end: Expr | int, + stride: Expr | int | None = None, ) -> Expr: """Returns a slice of the array. @@ -3537,9 +3598,7 @@ def array_slice( >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [[1, 2, 3, 4]]}) >>> result = df.select( - ... dfn.functions.array_slice( - ... dfn.col("a"), dfn.lit(2), dfn.lit(3) - ... ).alias("result")) + ... dfn.functions.array_slice(dfn.col("a"), 2, 3).alias("result")) >>> result.collect_column("result")[0].as_py() [2, 3] @@ -3547,18 +3606,27 @@ def array_slice( >>> result = df.select( ... dfn.functions.array_slice( - ... dfn.col("a"), dfn.lit(1), dfn.lit(4), - ... stride=dfn.lit(2), + ... dfn.col("a"), 1, 4, stride=2, ... ).alias("result")) >>> result.collect_column("result")[0].as_py() [1, 3] """ - if stride is not None: - stride = stride.expr - return Expr(f.array_slice(array.expr, begin.expr, end.expr, stride)) + begin = coerce_to_expr(begin) + end = coerce_to_expr(end) + stride = coerce_to_expr_or_none(stride) + return Expr( + f.array_slice( + array.expr, + begin.expr, + end.expr, + stride.expr if stride is not None else None, + ) + ) -def list_slice(array: Expr, begin: Expr, end: Expr, stride: Expr | None = None) -> Expr: +def list_slice( + array: Expr, begin: Expr | int, end: Expr | int, stride: Expr | int | None = None +) -> Expr: """Returns a slice of the array. See Also: @@ -3650,7 +3718,7 @@ def list_except(array1: Expr, array2: Expr) -> Expr: return array_except(array1, array2) -def array_resize(array: Expr, size: Expr, value: Expr) -> Expr: +def array_resize(array: Expr, size: Expr | int, value: Expr) -> Expr: """Returns an array with the specified size filled. If ``size`` is greater than the ``array`` length, the additional entries will @@ -3660,15 +3728,15 @@ def array_resize(array: Expr, size: Expr, value: Expr) -> Expr: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [[1, 2]]}) >>> result = df.select( - ... dfn.functions.array_resize(dfn.col("a"), dfn.lit(4), - ... dfn.lit(0)).alias("result")) + ... dfn.functions.array_resize(dfn.col("a"), 4, dfn.lit(0)).alias("result")) >>> result.collect_column("result")[0].as_py() [1, 2, 0, 0] """ + size = coerce_to_expr(size) return Expr(f.array_resize(array.expr, size.expr, value.expr)) -def list_resize(array: Expr, size: Expr, value: Expr) -> Expr: +def list_resize(array: Expr, size: Expr | int, value: Expr) -> Expr: """Returns an array with the specified size filled. If ``size`` is greater than the ``array`` length, the additional entries will be @@ -3822,7 +3890,7 @@ def list_zip(*arrays: Expr) -> Expr: def string_to_array( - string: Expr, delimiter: Expr, null_string: Expr | None = None + string: Expr, delimiter: Expr | str, null_string: Expr | str | None = None ) -> Expr: """Splits a string based on a delimiter and returns an array of parts. @@ -3832,9 +3900,7 @@ def string_to_array( >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": ["hello,world"]}) >>> result = df.select( - ... dfn.functions.string_to_array( - ... dfn.col("a"), dfn.lit(","), - ... ).alias("result")) + ... dfn.functions.string_to_array(dfn.col("a"), ",").alias("result")) >>> result.collect_column("result")[0].as_py() ['hello', 'world'] @@ -3842,17 +3908,24 @@ def string_to_array( >>> result = df.select( ... dfn.functions.string_to_array( - ... dfn.col("a"), dfn.lit(","), null_string=dfn.lit("world"), + ... dfn.col("a"), ",", null_string="world", ... ).alias("result")) >>> result.collect_column("result")[0].as_py() ['hello', None] """ - null_expr = null_string.expr if null_string is not None else None - return Expr(f.string_to_array(string.expr, delimiter.expr, null_expr)) + delimiter = coerce_to_expr(delimiter) + null_string = coerce_to_expr_or_none(null_string) + return Expr( + f.string_to_array( + string.expr, + delimiter.expr, + null_string.expr if null_string is not None else None, + ) + ) def string_to_list( - string: Expr, delimiter: Expr, null_string: Expr | None = None + string: Expr, delimiter: Expr | str, null_string: Expr | str | None = None ) -> Expr: """Splits a string based on a delimiter and returns an array of parts. diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 13c05a9e6..e0ebdbae5 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -964,12 +964,12 @@ def test_csv_read_options_builder_pattern(): options = ( CsvReadOptions() - .with_has_header(False) # noqa: FBT003 + .with_has_header(False) .with_delimiter("|") .with_quote("'") .with_schema_infer_max_records(2000) - .with_truncated_rows(True) # noqa: FBT003 - .with_newlines_in_values(True) # noqa: FBT003 + .with_truncated_rows(True) + .with_newlines_in_values(True) .with_file_extension(".tsv") ) assert options.has_header is False diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index 091fa9b56..9e2f791ea 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -3426,10 +3426,18 @@ def test_fill_null_all_null_column(ctx): assert result.column(1).to_pylist() == ["filled", "filled", "filled"] +_slow_udf_started = threading.Event() + + @udf([pa.int64()], pa.int64(), "immutable") def slow_udf(x: pa.Array) -> pa.Array: - # This must be longer than the check interval in wait_for_future - time.sleep(2.0) + _slow_udf_started.set() + # Sleep in small increments so Python's eval loop checks for pending + # async exceptions (like KeyboardInterrupt via PyThreadState_SetAsyncExc) + # between iterations. A single long time.sleep() is a C call where async + # exceptions are not checked on all Python versions (notably 3.11). + for _ in range(200): + time.sleep(0.01) return x @@ -3463,6 +3471,7 @@ def test_collect_or_stream_interrupted(slow_query, as_c_stream): # noqa: C901 P if as_c_stream: reader = pa.RecordBatchReader.from_stream(df) + _slow_udf_started.clear() read_started = threading.Event() read_exception = [] read_thread_id = None @@ -3474,6 +3483,14 @@ def trigger_interrupt(): msg = f"Read operation did not start within {max_wait_time} seconds" raise RuntimeError(msg) + # For slow_query tests, wait until the UDF is actually executing Python + # bytecode before sending the interrupt. PyThreadState_SetAsyncExc only + # delivers exceptions when the thread is in the Python eval loop, not + # while in native (Rust/C) code. + if slow_query and not _slow_udf_started.wait(timeout=max_wait_time): + msg = f"UDF did not start within {max_wait_time} seconds" + raise RuntimeError(msg) + if read_thread_id is None: msg = "Cannot get read thread ID" raise RuntimeError(msg) diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index d046eb48c..8aa791ae1 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -53,6 +53,8 @@ TransactionEnd, TransactionStart, Values, + coerce_to_expr, + coerce_to_expr_or_none, ensure_expr, ensure_expr_list, ) @@ -1030,12 +1032,55 @@ def test_ensure_expr_list_bytearray(): ensure_expr_list(bytearray(b"a")) +def test_coerce_to_expr_passes_expr_through(): + e = col("a") + result = coerce_to_expr(e) + assert isinstance(result, type(e)) + assert str(result) == str(e) + + +def test_coerce_to_expr_wraps_int(): + result = coerce_to_expr(42) + assert isinstance(result, type(lit(42))) + + +def test_coerce_to_expr_wraps_str(): + result = coerce_to_expr("hello") + assert isinstance(result, type(lit("hello"))) + + +def test_coerce_to_expr_wraps_float(): + result = coerce_to_expr(3.14) + assert isinstance(result, type(lit(3.14))) + + +def test_coerce_to_expr_wraps_bool(): + result = coerce_to_expr(True) + assert isinstance(result, type(lit(True))) + + +def test_coerce_to_expr_or_none_returns_none(): + assert coerce_to_expr_or_none(None) is None + + +def test_coerce_to_expr_or_none_wraps_value(): + result = coerce_to_expr_or_none(42) + assert isinstance(result, type(lit(42))) + + +def test_coerce_to_expr_or_none_passes_expr_through(): + e = col("a") + result = coerce_to_expr_or_none(e) + assert isinstance(result, type(e)) + assert str(result) == str(e) + + @pytest.mark.parametrize( "value", [ # Boolean - pa.scalar(True, type=pa.bool_()), # noqa: FBT003 - pa.scalar(False, type=pa.bool_()), # noqa: FBT003 + pa.scalar(True, type=pa.bool_()), + pa.scalar(False, type=pa.bool_()), # Integers - signed pa.scalar(127, type=pa.int8()), pa.scalar(-128, type=pa.int8()), diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 11e94af1c..d9781b1fb 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -2099,3 +2099,96 @@ def test_gen_series_with_step(): f.gen_series(literal(1), literal(10), literal(3)).alias("v") ).collect() assert result[0].column(0)[0].as_py() == [1, 4, 7, 10] + + +class TestPythonicNativeTypes: + """Tests for accepting native Python types instead of requiring lit().""" + + def test_split_part_native(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["a,b,c"]}) + result = df.select(f.split_part(column("a"), ",", 2).alias("s")).collect() + assert result[0].column(0)[0].as_py() == "b" + + def test_encode_native_str(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["hello"]}) + result = df.select(f.encode(column("a"), "base64").alias("e")).collect() + assert result[0].column(0)[0].as_py() == "aGVsbG8" + + def test_date_part_native_str(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["2021-07-15T00:00:00"]}) + df = df.select(f.to_timestamp(column("a")).alias("a")) + result = df.select(f.date_part("year", column("a")).alias("y")).collect() + assert result[0].column(0)[0].as_py() == 2021 + + def test_date_trunc_native_str(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["2021-07-15T12:34:56"]}) + df = df.select(f.to_timestamp(column("a")).alias("a")) + result = df.select(f.date_trunc("month", column("a")).alias("t")).collect() + assert str(result[0].column(0)[0].as_py()) == "2021-07-01 00:00:00" + + def test_left_native_int(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["the cat"]}) + result = df.select(f.left(column("a"), 3).alias("l")).collect() + assert result[0].column(0)[0].as_py() == "the" + + def test_round_native_int(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": [1.567]}) + result = df.select(f.round(column("a"), 2).alias("r")).collect() + assert result[0].column(0)[0].as_py() == 1.57 + + def test_regexp_count_native(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["abcabc"]}) + result = df.select( + f.regexp_count(column("a"), "abc", start=4, flags="i").alias("c") + ).collect() + assert result[0].column(0)[0].as_py() == 1 + + def test_log_native_int(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": [100.0]}) + result = df.select(f.log(10, column("a")).alias("l")).collect() + assert result[0].column(0)[0].as_py() == 2.0 + + def test_power_native_int(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": [2.0]}) + result = df.select(f.power(column("a"), 3).alias("p")).collect() + assert result[0].column(0)[0].as_py() == 8.0 + + def test_array_slice_native(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3, 4]]}) + result = df.select(f.array_slice(column("a"), 2, 3).alias("s")).collect() + assert result[0].column(0)[0].as_py() == [2, 3] + + def test_string_to_array_native(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["hello,NA,world"]}) + result = df.select( + f.string_to_array(column("a"), ",", null_string="NA").alias("v") + ).collect() + assert result[0].column(0)[0].as_py() == ["hello", None, "world"] + + def test_regexp_replace_native(self): + ctx = SessionContext() + df = ctx.from_pydict({"a": ["a1 b2 c3"]}) + result = df.select( + f.regexp_replace(column("a"), r"\d+", "X", flags="g").alias("r") + ).collect() + assert result[0].column(0)[0].as_py() == "aX bX cX" + + def test_backward_compat_with_lit(self): + """Verify that existing code using lit() still works.""" + ctx = SessionContext() + df = ctx.from_pydict({"a": ["a,b,c"]}) + result = df.select( + f.split_part(column("a"), literal(","), literal(2)).alias("s") + ).collect() + assert result[0].column(0)[0].as_py() == "b" From c657dad97349c1113e843e3e15bb41f865e65a97 Mon Sep 17 00:00:00 2001 From: Nick <24689722+ntjohnson1@users.noreply.github.com> Date: Wed, 29 Apr 2026 07:35:21 -0400 Subject: [PATCH 30/83] Move public skills to a directory to avoid downloading the whole repo (#1519) * Move to skill directory * Avoid moved skill with test --- dev/release/rat_exclude_files.txt | 2 +- SKILL.md => skills/datafusion_python/SKILL.md | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename SKILL.md => skills/datafusion_python/SKILL.md (100%) diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index a7a497dab..6769a245c 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -49,4 +49,4 @@ 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 \ No newline at end of file diff --git a/SKILL.md b/skills/datafusion_python/SKILL.md similarity index 100% rename from SKILL.md rename to skills/datafusion_python/SKILL.md From 13b2c47b0d5e348cea24b9264e87fd67666c56f1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sun, 3 May 2026 10:52:10 -0400 Subject: [PATCH 31/83] Update user documentation for AI agent skill usage (#1505) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) * 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) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- 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 +++++ .../user-guide/ai-coding-assistants.rst | 82 +++++++++++ .../common-operations/aggregations.rst | 50 +++++++ .../common-operations/expressions.rst | 92 +++++++++++++ .../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(-) create mode 100644 docs/source/llms.txt create mode 100644 docs/source/user-guide/ai-coding-assistants.rst create mode 100644 examples/array-operations.py diff --git a/AGENTS.md b/AGENTS.md index 7d3262710..632d6ebc0 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 4baed7d1d..f6ee662d0 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 6769a245c..d93ff3c24 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 01813b032..b2e9bb8c3 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -129,6 +129,10 @@ def setup(sphinx) -> None: # 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 134d41cb6..0007cc41a 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 000000000..3ff6b3813 --- /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 000000000..fb7998c6d --- /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 `_ — + 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 `_ open standard. +- `llms.txt `_ — an entry point for LLM-based tools following the + `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 de24a2ba5..f59b62ab4 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 7848b4ee7..ae1ccc0dc 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 f669721a3..59c47b595 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 `_ 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 0ef194afe..3024c782f 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 000000000..884f93974 --- /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 e4972411a..f08b464bb 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -52,7 +52,7 @@ 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 From db22a9240e5832d9a88b74c640e3b60abfbe52c1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 6 May 2026 15:06:53 -0400 Subject: [PATCH 32/83] docs: add upstream sync process documentation (#1524) * docs: add upstream sync process documentation Document the three-PR workflow used to sync to a newer upstream apache/datafusion version: bump crate deps + fix breakage, consolidate transitive deps, then fill API and documentation gaps via /check-upstream. Cross-reference from dev/release/README.md. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: add audit-skill-md skill New AI agent skill at .ai/skills/audit-skill-md/SKILL.md to keep the user-facing skills/datafusion_python/SKILL.md in sync with the public Python API. Audits SessionContext, DataFrame, Expr, and functions surfaces for new APIs not covered, stale mentions, examples that drifted from idiomatic style, and missing version notes. Wired into PR 3 of the upstream sync workflow documented in dev/release/upstream-sync.md. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: verify upstream sync completed before release Add a checklist item to "Preparing the main Branch" pointing release managers at dev/release/upstream-sync.md so the crate bump, dependency consolidation, and /check-upstream and /audit-skill-md passes are confirmed done before the release branch is cut. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: scope upstream sync cargo update to datafusion family Replace `cargo update -p datafusion` with an explicit multi-`-p` invocation listing every `datafusion-*` workspace dependency, so PR 1 of the upstream-sync workflow refreshes only the datafusion family and leaves other transitives for PR 2 to consolidate. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: correct datafusion-* pin location in upstream sync PR 1 step 1 incorrectly stated downstream `datafusion-*` crates are pinned in `crates/core/Cargo.toml`. Pins live in the root `[workspace.dependencies]`; per-crate manifests inherit via `workspace = true`. Reword step 1 to point at the right file. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: restore workspace.package version bump in upstream sync PR 1 step 1 must also bump `[workspace.package].version` because the `datafusion-python` major version tracks the upstream `datafusion` major. The previous reword dropped that instruction. Reinstate it alongside the `[workspace.dependencies]` updates. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: align audit-skill-md description with body version phrasing Frontmatter description referenced "requires upstream DataFusion vX", but the body of the skill settles on the `datafusion-python NN` form (consistent with the package/upstream-major equivalence). Switch the description to match so the skill speaks one language end to end. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: fold make-pythonic step into PR 3 of upstream sync Audit-skill-md documents the order `/check-upstream` -> `/make-pythonic` (optional) -> `/audit-skill-md`, but PR 3 of the upstream-sync workflow only listed the first and last. Insert the make-pythonic pass as step 3 so signatures get aligned before the SKILL.md audit, avoiding example churn. Drops the orphan trailing paragraph in favor of inline guidance on when to defer larger reshapes to their own PR. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: drop literal Cargo.toml version from audit-skill-md inputs Replace literal `version = "53.0.0"` example with a pointer to the `[workspace.package]` field plus an `NN.0.0` placeholder so the skill prose does not drift each major bump. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .ai/skills/audit-skill-md/SKILL.md | 284 +++++++++++++++++++++++++++++ dev/release/README.md | 9 + dev/release/upstream-sync.md | 170 +++++++++++++++++ 3 files changed, 463 insertions(+) create mode 100644 .ai/skills/audit-skill-md/SKILL.md create mode 100644 dev/release/upstream-sync.md diff --git a/.ai/skills/audit-skill-md/SKILL.md b/.ai/skills/audit-skill-md/SKILL.md new file mode 100644 index 000000000..30e1a90fd --- /dev/null +++ b/.ai/skills/audit-skill-md/SKILL.md @@ -0,0 +1,284 @@ + + +--- +name: audit-skill-md +description: Audit the user-facing skill at skills/datafusion_python/SKILL.md against the current public Python API. Find new APIs that should be documented, stale mentions of removed/renamed APIs, examples that drifted from current idiomatic style, and places that need a "requires datafusion-python NN or newer" note. Run after upstream syncs and before each release. +argument-hint: [scope] (e.g., "session-context", "dataframe", "expr", "functions", "patterns", "pitfalls", "version-notes", "all") +--- + +# Audit `skills/datafusion_python/SKILL.md` + +You are auditing the user-facing skill at +[`skills/datafusion_python/SKILL.md`](../../skills/datafusion_python/SKILL.md) +against the current state of the Python API. The skill is the source of truth +for how AI coding assistants are taught to write `datafusion-python` code, so +it must match what the project actually ships. This skill identifies gaps +caused by upstream syncs, refactors, or renames, and (if asked) applies the +edits directly to `SKILL.md`. + +The skill is most usefully run **after** the `check-upstream` step of an +upstream sync (see `dev/release/upstream-sync.md`) — once any new APIs are +exposed, this skill makes sure they get documented. + +## What the skill covers + +The user-facing `SKILL.md` documents these public surfaces. This list is not +exhaustive — if a new top-level area is added (e.g., a new `Catalog` API +exposed at the package root), include it. + +| Surface | Module | Sections in SKILL.md | +|---|---|---| +| `SessionContext` | `python/datafusion/context.py` | "Data Loading" | +| `DataFrame` | `python/datafusion/dataframe.py` | "DataFrame Operations Quick Reference", "Executing and Collecting Results", "Idiomatic Patterns" | +| `Expr` | `python/datafusion/expr.py` | "Expression Building", "Common Pitfalls" | +| `functions` | `python/datafusion/functions.py` | "Available Functions (Categorized)", scattered uses throughout | +| Top-level helpers (`col`, `lit`, `WindowFrame`, ...) | `python/datafusion/__init__.py` | "Import Conventions", "Core Abstractions" | + +## Scope argument + +The user may specify a scope via `$ARGUMENTS` to limit the audit. If no scope +is given or `all` is specified, audit every area. + +| Scope | Audit target | +|---|---| +| `session-context` | `SessionContext` methods and the "Data Loading" section | +| `dataframe` | `DataFrame` methods and the operations / executing / patterns sections | +| `expr` | `Expr` methods/operators and the "Expression Building" section | +| `functions` | `functions.py` `__all__` and the "Available Functions (Categorized)" section | +| `patterns` | "Idiomatic Patterns" section — confirm patterns still match recommended style | +| `pitfalls` | "Common Pitfalls" — confirm each pitfall still reproduces, drop ones fixed upstream | +| `version-notes` | Cross-check version annotations (see below) | +| `all` | Everything above | + +## Inputs to read + +Before producing the report: + +1. `skills/datafusion_python/SKILL.md` — the document being audited. +2. The relevant Python module(s) for the chosen scope. Public surface is the + `__all__` list (where defined) plus `class` and `def` symbols not prefixed + with `_`. +3. `Cargo.toml` (root) for the current `datafusion-python` version — read + the `version` field under `[workspace.package]` (format `NN.0.0`). The + major version always matches the upstream `datafusion` crate, so a + single `datafusion-python` version expresses both. + `python/datafusion/__init__.py`'s `__version__` is the same value + exposed at runtime. +4. Recent commits touching the relevant module(s) for context on what + changed since the last sync: + ```bash + git log --oneline -- python/datafusion/dataframe.py | head -20 + ``` + +## What to look for + +Walk through each scoped area and flag four kinds of issues. + +### 1. New APIs not mentioned + +For each public symbol in the module's `__all__` (or each public class +method), check whether it appears anywhere in `SKILL.md`. A symbol is +"covered" if it shows up in: + +- A code block (the strongest signal — it's demonstrated). +- The "Available Functions (Categorized)" list. +- The SQL-to-DataFrame Reference table. + +**Decide whether each missing symbol deserves an entry.** Not every public +symbol belongs in `SKILL.md` — the skill is curated for the patterns users +hit daily, not exhaustive API reference. Use these heuristics: + +- **Add it** if it replaces or supersedes something already in the skill + (e.g., a new operation that is the idiomatic alternative to a documented + workaround). +- **Add it** if it fits a category already present (a new aggregate function + goes in the aggregate list; a new join type goes in the joining section). +- **Add it** if it changes how a documented pattern should be written. +- **Skip it** if it is genuinely niche / advanced / experimental. +- **Skip it** if it is internal plumbing exposed for FFI but not user-facing. + +When you flag a missing symbol, include a one-line proposed insertion point +(which section / which table row) so a reviewer can decide quickly. + +### 2. Stale mentions + +For each function name, method name, or import shown in `SKILL.md`, verify it +still exists in the current API: + +- Function names mentioned in prose or in the categorized list should appear + in `python/datafusion/functions.py`'s `__all__`. +- Method calls in code blocks should resolve against the current class. +- Imports (`from datafusion import ...`) should succeed against the current + `__init__.py`. + +A quick way to check imports without running them: + +```bash +python -c "from datafusion import SessionContext, col, lit; from datafusion import functions as F; print('ok')" +``` + +For each stale mention, propose either: +- a rename to the current name, or +- removal if the API is gone with no replacement. + +### 3. Examples that drifted from idiomatic style + +The skill teaches a Pythonic style: prefer plain strings to `col(...)` when a +column reference is all you need; prefer raw Python values to `lit(...)` +where auto-wrapping applies. Recent refactors (see the `make-pythonic` +skill) keep moving more functions toward accepting native types. + +For each code example in `SKILL.md`, check: + +- Does it use `lit(value)` where a raw value would work? Comparison RHS, + arithmetic with a column, etc. all auto-wrap. (Reserve `lit()` for the + cases listed in pitfall #2.) +- Does it use `col("name")` where a plain string would work? `select(...)`, + `aggregate([keys], ...)`, `sort(...)`, `sort_by(...)` all accept plain + name strings. +- Do `functions.py` calls match the current pythonic signature for that + function? If `make-pythonic` recently changed a signature (e.g., + `repeat(string, n: Expr | int)`), the example should pass `3` rather than + `lit(3)`. +- Does any example use a deprecated or removed parameter name? + +For drift, propose the updated snippet. If the change is purely stylistic +and the older form still works, mark the suggestion as **non-blocking**. + +### 4. Missing or stale version notes + +When an API depends on a specific version, the skill should say so — +otherwise an agent referencing the skill in an older project will write +code that fails at import or at runtime. + +`datafusion-python` shares its major version number with the upstream +`datafusion` crate (e.g., `datafusion-python 53.x` tracks upstream +`datafusion 53`). Always express version requirements in terms of +`datafusion-python` only — there is no need to call out upstream and +package versions separately. + +Add a version note when: + +- A method or function shown in the skill was added in a specific release + (e.g., a new `DataFrame` method that didn't exist before 53). +- A breaking change altered behavior in a specific release (signature + change, default-value change, new required argument). +- A pitfall was fixed in a specific release. Either annotate the pitfall + block with "fixed in datafusion-python NN, kept here for users on older + versions" or remove it once the supported floor moves past that version. + +Format for version notes (inline, italicized): + +```markdown +*Requires datafusion-python 53 or newer.* +``` + +For each missing/stale version note, propose the exact line and where it +belongs. + +## How to discover changes since the last audit + +If the user supplies a previous version or commit SHA where the audit was +last run, diff against it: + +```bash +# Public-API-relevant changes since SHA +git log --oneline ..HEAD -- python/datafusion/ + +# Whose signatures actually moved +git diff ..HEAD -- python/datafusion/functions.py | grep '^[+-]def ' +``` + +If no prior audit point is given, fall back to "since the last upstream +sync" by inspecting commits that touch `Cargo.toml`'s `datafusion` pin: + +```bash +git log --oneline -- Cargo.toml | grep -i datafusion | head -5 +``` + +## Output Format + +Produce a report grouped by scope. Each finding is one bullet with a +proposed action, so a maintainer can review the list quickly and apply +edits in order. + +``` +## SKILL.md Audit (scope: ) + +Audited against: +- skills/datafusion_python/SKILL.md @ +- datafusion-python + +### New APIs to cover +- `DataFrame.foo()` — added in datafusion-python 53. Insert in "DataFrame Operations Quick Reference" under . + Proposed snippet: + ```python + df.foo(...) + ``` + +### Stale mentions +- "old_function_name" referenced in the categorized list (line N) — renamed to "new_function_name". Replace. + +### Drifted examples +- "Filtering" section, `df.filter(col("a") > lit(10))` — drop `lit(10)`, auto-wrap applies. (non-blocking) +- "Aggregation" section, `df.aggregate([col("region")], ...)` — pass `"region"` as a plain string per "Projection" guidance. + +### Version notes +- `DataFrame.foo()` block needs *Requires datafusion-python 53 or newer.* +- "Common Pitfalls" #N — fixed in datafusion-python 53; remove the pitfall and update the SQL-to-DataFrame row to no longer flag the workaround. + +### No-change confirmed +- `SessionContext` data-loading section — all entries match current API. +``` + +If asked to apply the changes, edit `skills/datafusion_python/SKILL.md` +directly with `Edit` tool calls, one finding at a time, and re-run the +relevant doctest sanity check at the end: + +```bash +pytest --doctest-modules python/datafusion -q +``` + +## What NOT to flag + +- **Internal helpers / underscored names.** Private symbols are not part of + the user-facing surface. +- **Functions intentionally omitted.** Niche / advanced APIs (custom + catalogs, raw FFI plumbing, low-level execution plan accessors) live in + the API reference, not the skill. If an omission was deliberate and a + comment / commit explains why, leave it out. +- **Style nits inside explanatory prose.** The skill mixes example code and + prose; only enforce the pythonic style on actual code blocks. +- **Function-by-function coverage of every `functions.py` symbol.** The + "Available Functions (Categorized)" list is curated by category, not + exhaustive. Adding a single new aggregate to the aggregate list is + enough — the user follows the pointer to the API reference for the rest. + +## Coordination with other skills + +- Run `/check-upstream` first to expose any missing upstream APIs into the + Python layer. Without that, this skill cannot recommend documenting + something that is not yet exposed. +- Run `/make-pythonic` before this skill if a Pythonic-signature pass is + planned for a release — that way this skill can update examples to the + final signature in one shot rather than churning them twice. +- The order during an upstream sync (PR 3 of `dev/release/upstream-sync.md`) + is therefore: `/check-upstream` → `/make-pythonic` (optional) → + `/audit-skill-md`. diff --git a/dev/release/README.md b/dev/release/README.md index 4833be55a..1b02ca576 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -33,6 +33,12 @@ release branch without blocking ongoing development in the `main` branch. We can cherry-pick commits from the `main` branch into `branch-53` as needed and then create new patch releases from that branch. +## Upstream Sync + +Between releases the `main` branch is periodically synced to a newer upstream `apache/datafusion` version. This is +broken into a three-PR workflow (bump + fix breakage, consolidate transitive deps, fill API and documentation gaps). +See [`upstream-sync.md`](upstream-sync.md) for the full process. + ## Detailed Guide ### Pre-requisites @@ -53,6 +59,9 @@ You will also need access to the [datafusion](https://test.pypi.org/project/data Before creating a new release: - We need to ensure that the main branch does not have any GitHub dependencies +- Confirm the upstream sync workflow in [`upstream-sync.md`](upstream-sync.md) has been completed for this release cycle + (crate bump + breakage fixes, transitive dependency consolidation, and the `/check-upstream` and `/audit-skill-md` + passes). Any gaps surfaced by those skills should land before the release branch is cut. - a PR should be created and merged to update the major version number of the project - A new release branch should be created, such as `branch-53` - It is best to push this branch to the apache repository rather than a personal fork in case patch releases are required. diff --git a/dev/release/upstream-sync.md b/dev/release/upstream-sync.md new file mode 100644 index 000000000..dc05b4195 --- /dev/null +++ b/dev/release/upstream-sync.md @@ -0,0 +1,170 @@ + + +# Upstream Sync Process + +This document describes how to sync `datafusion-python` to a new version of the +upstream `apache/datafusion` Rust crates. This is a recurring task: between +official releases the `main` branch tracks DataFusion via crates.io or GitHub +dependencies, and we periodically bump those dependencies to pick up new +features and bug fixes. + +The work is broken into **three sequential PRs** rather than landing as one +large change. Splitting reviews along these lines keeps each PR focused, makes +breakage easier to bisect, and lets reviewers concentrate on one concern at a +time. + +## PR 1: Bump DataFusion crate dependencies and fix breakage + +**Goal:** update the upstream `datafusion` crate version and make the project +build, test, and lint cleanly against it. + +1. In the root `Cargo.toml`, update: + - `[workspace.package].version` to the new major (the `datafusion-python` + major tracks the upstream `datafusion` major, so a 53→54 bump moves + this from `53.0.0` to `54.0.0`), and + - every `datafusion` / `datafusion-*` entry in `[workspace.dependencies]` + to the same new major. + + Per-crate manifests under `crates/` inherit these pins via + `workspace = true` and need no edit. +2. Update `Cargo.lock` for the datafusion family only — leave unrelated + transitives at their current pins so PR 2 can address them deliberately. + List every `datafusion-*` workspace dependency with `-p`: + + ```bash + cargo update \ + -p datafusion \ + -p datafusion-substrait \ + -p datafusion-proto \ + -p datafusion-ffi \ + -p datafusion-catalog \ + -p datafusion-common \ + -p datafusion-functions-aggregate \ + -p datafusion-functions-window \ + -p datafusion-expr + ``` + + Or pin exact versions with `--precise`, one crate at a time: + + ```bash + cargo update -p datafusion --precise 54.0.0 + # repeat for each datafusion-* sibling + ``` + + A bare `cargo update` would refresh every transitive crate and blur the + diff between PR 1 and PR 2. +3. Run the standard build and test commands and address compilation errors, + API renames, signature changes, and behavior changes: + - `cargo build` + - `cargo test` + - `pytest` + - `pre-commit run --all-files` +4. Fix only what's needed to restore green CI. Resist the urge to bundle + unrelated cleanups — those belong in their own PR. +5. If a breaking change in upstream requires a user-facing API change in + `datafusion-python`, add the `api change` label and document the change + in the PR description so it surfaces in the changelog. + +**Reference PRs:** [#1311](https://github.com/apache/datafusion-python/pull/1311) +(DF51), [#1337](https://github.com/apache/datafusion-python/pull/1337) (DF52). + +## PR 2: Consolidate transitive dependencies + +**Goal:** after the upstream bump, the dependency tree may have multiple +versions of the same transitive crate (for example, two `arrow` versions, two +`object_store` versions). Reconcile these so we ship a single coherent set. + +1. Inspect the lockfile for duplicates: + ```bash + cargo tree --duplicates + ``` +2. For each duplicate that matters (Arrow, `object_store`, `parquet`, + `tokio`, `arrow-flight`, etc.), update our direct dependency declarations + in `Cargo.toml` to versions compatible with what upstream DataFusion now + pulls in. The goal is one version of each ecosystem-critical crate. +3. Re-run `cargo update` and re-run the full test matrix. Some duplicates are + benign (small leaf crates with no FFI surface) and can be left alone if + reconciliation would force a much larger change. Use judgment. +4. If consolidating forces a behavioral change visible to users (for example, + a newer `pyarrow`-compatible Arrow version), call it out in the PR + description. + +Keeping this work separate from PR 1 means PR 1 stays a "make it compile" +review and PR 2 stays a "tidy the dependency graph" review. + +## PR 3: Fill API and documentation gaps + +**Goal:** with the upstream version locked in, identify new APIs that landed +upstream and decide whether to expose them, and update agent-facing +documentation so it still matches the surface we ship. + +1. Run the `check-upstream` skill (`.ai/skills/check-upstream/SKILL.md`) to + diff the upstream Rust API against what's exposed in + `python/datafusion/`. The skill covers scalar/aggregate/window/table + functions, `DataFrame` methods, `SessionContext` methods, and FFI types. + Invoke it from the assistant with `/check-upstream` (optionally scoped to + one area, e.g. `/check-upstream scalar functions`). +2. For each gap, decide whether to: + - Expose it now (small, obvious additions can land in this PR). + - File a tracking issue (anything non-trivial — separate PR per feature + keeps reviews focused). + - Skip it (internal-only or already covered by an existing API; record + the decision in the "Evaluated and not requiring exposure" sections of + the skill so future runs don't re-flag it). +3. (Optional) Run the `make-pythonic` skill + (`.ai/skills/make-pythonic/SKILL.md`) over any newly exposed APIs to + align signatures with the project's Pythonic style (accepting plain + strings for column names, raw Python values where auto-wrapping + applies, etc.). Invoke it from the assistant with `/make-pythonic`. + Running this *before* the audit step means examples in `SKILL.md` get + updated to the final signature in one pass instead of churning twice. + Larger reshapes still belong in their own PR. +4. Run the `audit-skill-md` skill (`.ai/skills/audit-skill-md/SKILL.md`) to + cross-reference the user-facing skill at + [`skills/datafusion_python/SKILL.md`](../../skills/datafusion_python/SKILL.md) + against the current public API. The skill flags stale function names, + missing newly exposed APIs, examples that drifted from idiomatic style, + and missing version notes. Invoke it from the assistant with + `/audit-skill-md` (optionally scoped, e.g. `/audit-skill-md dataframe`). + Apply the resulting edits to `SKILL.md` and to the relevant RST pages + under `docs/source/user-guide/common-operations/`. +5. If new aggregate or window functions were exposed in step 2, also update: + - `docs/source/user-guide/common-operations/aggregations.rst` + - `docs/source/user-guide/common-operations/windows.rst` + +## Why three PRs + +- **Bisectable.** If a regression appears, `git bisect` lands on the + responsible PR (compile fix, dependency consolidation, or API addition) + rather than a single mega-commit. +- **Reviewable.** Each PR has a single concern. Reviewers reading PR 1 don't + need to also reason about whether new APIs are well-named. +- **Skippable.** Some upstream syncs are pure version bumps with no new APIs + worth exposing. PR 3 can be empty or merged as a no-op if the audit comes + back clean. + +## Related documents + +- [`README.md`](README.md) — the broader release process (this sync work + feeds into the next official release). +- [`.ai/skills/check-upstream/SKILL.md`](../../.ai/skills/check-upstream/SKILL.md) + — API coverage audit. +- [`skills/datafusion_python/SKILL.md`](../../skills/datafusion_python/SKILL.md) + — user-facing agent guide kept in sync via PR 3. From 55870ff30c4a1f086e0b63de434ebf8b674ac110 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 13 May 2026 10:28:12 -0400 Subject: [PATCH 33/83] build(deps): combined dependabot bumps (Cargo + workflows) (#1534) Combines the following dependabot PRs: Cargo: - tokio 1.50.0 -> 1.52.3 (#1530) - arrow-array 58.1.0 -> 58.3.0 (#1522, also picks up 58.3.0) - arrow-schema 58.1.0 -> 58.3.0 (#1523, also picks up 58.3.0) - datafusion 53.0.0 -> 53.1.0 (#1515) - datafusion-catalog 53.0.0 -> 53.1.0 (#1513) - datafusion-ffi 53.0.0 -> 53.1.0 (#1512) - datafusion-proto 53.0.0 -> 53.1.0 (#1511) - datafusion-common 53.0.0 -> 53.1.0 (#1510) - mimalloc 0.1.48 -> 0.1.50 (#1514) - uuid 1.23.0 -> 1.23.1 (#1508) - rustls-webpki 0.103.10 -> 0.103.13 (#1506) - rand 0.9.2 -> 0.9.4 (#1495) GitHub Actions: - github/codeql-action 4.32.5 -> 4.35.4 (#1531) - astral-sh/setup-uv 7.3.1 -> 8.1.0 (#1500) Each individual PR was passing CI before this combined bump. Co-authored-by: Claude Opus 4.7 (1M context) --- .github/workflows/build.yml | 14 +-- .github/workflows/codeql.yml | 4 +- .github/workflows/test.yml | 2 +- Cargo.lock | 207 ++++++++++++++++++----------------- Cargo.toml | 2 +- 5 files changed, 117 insertions(+), 112 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7682d6cb0..37a9dba03 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -69,7 +69,7 @@ jobs: with: python-version: "3.12" - - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true @@ -113,7 +113,7 @@ jobs: steps: - uses: actions/checkout@v6 - - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true @@ -155,7 +155,7 @@ jobs: with: key: ${{ inputs.build_mode }} - - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true @@ -236,7 +236,7 @@ jobs: with: key: ${{ inputs.build_mode }} - - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true @@ -307,7 +307,7 @@ jobs: with: key: ${{ inputs.build_mode }} - - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true @@ -381,7 +381,7 @@ jobs: with: key: ${{ inputs.build_mode }} - - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true @@ -505,7 +505,7 @@ jobs: python-version: "3.10" - name: Install dependencies - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index a9855cf48..e71ea6bed 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -44,11 +44,11 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@c793b717bc78562f491db7b0e93a3a178b099162 # v4 + uses: github/codeql-action/init@68bde559dea0fdcac2102bfdf6230c5f70eb485e # v4 with: languages: actions - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@c793b717bc78562f491db7b0e93a3a178b099162 # v4 + uses: github/codeql-action/analyze@68bde559dea0fdcac2102bfdf6230c5f70eb485e # v4 with: category: "/language:actions" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 706ccbc55..c597ab308 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -56,7 +56,7 @@ jobs: key: cargo-cache-${{ matrix.toolchain }}-${{ hashFiles('Cargo.lock') }} - name: Install dependencies - uses: astral-sh/setup-uv@5a095e7a2014a4212f075830d4f7277575a9d098 + uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true diff --git a/Cargo.lock b/Cargo.lock index 4efca3eb6..70f09ec46 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -131,7 +131,7 @@ dependencies = [ "miniz_oxide", "num-bigint", "quad-rand", - "rand 0.9.2", + "rand 0.9.4", "regex-lite", "serde", "serde_bytes", @@ -212,9 +212,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "772bd34cacdda8baec9418d80d23d0fb4d50ef0735685bd45158b83dfeb6e62d" +checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" dependencies = [ "ahash", "arrow-buffer", @@ -223,7 +223,7 @@ dependencies = [ "chrono", "chrono-tz", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "num-complex", "num-integer", "num-traits", @@ -231,9 +231,9 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "898f4cf1e9598fdb77f356fdf2134feedfd0ee8d5a4e0a5f573e7d0aec16baa4" +checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" dependencies = [ "bytes", "half", @@ -280,9 +280,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42d10beeab2b1c3bb0b53a00f7c944a178b622173a5c7bcabc3cb45d90238df4" +checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" dependencies = [ "arrow-buffer", "arrow-schema", @@ -371,9 +371,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c30a1365d7a7dc50cc847e54154e6af49e4c4b0fddc9f607b687f29212082743" +checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" dependencies = [ "bitflags", "serde_core", @@ -920,9 +920,9 @@ dependencies = [ [[package]] name = "datafusion" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de9f8117889ba9503440f1dd79ebab32ba52ccf1720bb83cd718a29d4edc0d16" +checksum = "93db0e623840612f7f2cd757f7e8a8922064192363732c88692e0870016e141b" dependencies = [ "arrow", "arrow-schema", @@ -964,7 +964,7 @@ dependencies = [ "object_store", "parking_lot", "parquet", - "rand 0.9.2", + "rand 0.9.4", "regex", "sqlparser", "tempfile", @@ -976,9 +976,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be893b73a13671f310ffcc8da2c546b81efcc54c22e0382c0a28aa3537017137" +checksum = "37cefde60b26a7f4ff61e9d2ff2833322f91df2b568d7238afe67bde5bdffb66" dependencies = [ "arrow", "async-trait", @@ -1001,9 +1001,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830487b51ed83807d6b32d6325f349c3144ae0c9bf772cf2a712db180c31d5e6" +checksum = "17e112307715d6a7a331111a4c2330ff54bc237183511c319e3708a4cff431fb" dependencies = [ "arrow", "async-trait", @@ -1024,9 +1024,9 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d7663f3af955292f8004e74bcaf8f7ea3d66cc38438749615bb84815b61a293" +checksum = "d72a11ca44a95e1081870d3abb80c717496e8a7acb467a1d3e932bb636af5cc2" dependencies = [ "ahash", "apache-avro", @@ -1050,9 +1050,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f590205c7e32fe1fea48dd53ffb406e56ae0e7a062213a3ac848db8771641bd" +checksum = "89f4afaed29670ec4fd6053643adc749fe3f4bc9d1ce1b8c5679b22c67d12def" dependencies = [ "futures", "log", @@ -1061,9 +1061,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fde1e030a9dc87b743c806fbd631f5ecfa2ccaa4ffb61fa19144a07fea406b79" +checksum = "e9fb386e1691355355a96419978a0022b7947b44d4a24a6ea99f00b6b485cbb6" dependencies = [ "arrow", "async-compression", @@ -1087,7 +1087,7 @@ dependencies = [ "liblzma", "log", "object_store", - "rand 0.9.2", + "rand 0.9.4", "tokio", "tokio-util", "url", @@ -1096,9 +1096,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "331ebae7055dc108f9b54994b93dff91f3a17445539efe5b74e89264f7b36e15" +checksum = "ffa6c52cfed0734c5f93754d1c0175f558175248bf686c944fb05c373e5fc096" dependencies = [ "arrow", "arrow-ipc", @@ -1120,9 +1120,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49dda81c79b6ba57b1853a9158abc66eb85a3aa1cede0c517dabec6d8a4ed3aa" +checksum = "a579c3bd290c66ea4b269493e75e8a3ed42c9c895a651f10210a29538aee50c4" dependencies = [ "apache-avro", "arrow", @@ -1140,9 +1140,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e0d475088325e2986876aa27bb30d0574f72a22955a527d202f454681d55c5c" +checksum = "503f29e0582c1fc189578d665ff57d9300da1f80c282777d7eb67bb79fb8cdca" dependencies = [ "arrow", "async-trait", @@ -1163,9 +1163,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea1520d81f31770f3ad6ee98b391e75e87a68a5bb90de70064ace5e0a7182fe8" +checksum = "e33804749abc8d0c8cb7473228483cb8070e524c6f6086ee1b85a64debe2b3d2" dependencies = [ "arrow", "async-trait", @@ -1187,9 +1187,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95be805d0742ab129720f4c51ad9242cd872599cdb076098b03f061fcdc7f946" +checksum = "32a8e0365e0e08e8ff94d912f0ababcf9065a1a304018ba90b1fc83c855b4997" dependencies = [ "arrow", "async-trait", @@ -1217,15 +1217,15 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c93ad9e37730d2c7196e68616f3f2dd3b04c892e03acd3a8eeca6e177f3c06a" +checksum = "8de6ac0df1662b9148ad3c987978b32cbec7c772f199b1d53520c8fa764a87ee" [[package]] name = "datafusion-execution" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9437d3cd5d363f9319f8122182d4d233427de79c7eb748f23054c9aaa0fdd8df" +checksum = "c03c7fbdaefcca4ef6ffe425a5fc2325763bfb426599bb0bf4536466efabe709" dependencies = [ "arrow", "arrow-buffer", @@ -1239,16 +1239,16 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.2", + "rand 0.9.4", "tempfile", "url", ] [[package]] name = "datafusion-expr" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67164333342b86521d6d93fa54081ee39839894fb10f7a700c099af96d7552cf" +checksum = "574b9b6977fedbd2a611cbff12e5caf90f31640ad9dc5870f152836d94bad0dd" dependencies = [ "arrow", "async-trait", @@ -1269,9 +1269,9 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab05fdd00e05d5a6ee362882546d29d6d3df43a6c55355164a7fbee12d163bc9" +checksum = "7d7c3adf3db8bf61e92eb90cb659c8e8b734593a8f7c8e12a843c7ddba24b87e" dependencies = [ "arrow", "datafusion-common", @@ -1282,9 +1282,9 @@ dependencies = [ [[package]] name = "datafusion-ffi" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b8250f7cdf463a0ad145f41d7508bcfa54c9b9f027317e599f0331097e3cc38" +checksum = "b95173344d04ba62755c949bf44f8d1a6e4414cf6392a635db96c07e711b9a3c" dependencies = [ "abi_stable", "arrow", @@ -1332,9 +1332,9 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fb863482d987cf938db2079e07ab0d3bb64595f28907a6c2f8671ad71cca7e" +checksum = "f28aa4e10384e782774b10e72aca4d93ef7b31aa653095d9d4536b0a3dbc51b6" dependencies = [ "arrow", "arrow-buffer", @@ -1355,7 +1355,7 @@ dependencies = [ "md-5", "memchr", "num-traits", - "rand 0.9.2", + "rand 0.9.4", "regex", "sha2", "unicode-segmentation", @@ -1364,9 +1364,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "829856f4e14275fb376c104f27cbf3c3b57a9cfe24885d98677525f5e43ce8d6" +checksum = "00aa6217e56098ba84e0a338176fe52f0a84cca398021512c6c8c5eff806d0ad" dependencies = [ "ahash", "arrow", @@ -1386,9 +1386,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08af79cc3d2aa874a362fb97decfcbd73d687190cb096f16a6c85a7780cce311" +checksum = "b511250349407db7c43832ab2de63f5557b19a20dfd236b39ca2c04468b50d47" dependencies = [ "ahash", "arrow", @@ -1399,9 +1399,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "465ae3368146d49c2eda3e2c0ef114424c87e8a6b509ab34c1026ace6497e790" +checksum = "ef13a858e20d50f0a9bb5e96e7ac82b4e7597f247515bccca4fdd2992df0212a" dependencies = [ "arrow", "arrow-ord", @@ -1424,9 +1424,9 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6156e6b22fcf1784112fc0173f3ae6e78c8fdb4d3ed0eace9543873b437e2af6" +checksum = "72b40d3f5bbb3905f9ccb1ce9485a9595c77b69758a7c24d3ba79e334ff51e7e" dependencies = [ "arrow", "async-trait", @@ -1440,9 +1440,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca7baec14f866729012efb89011a6973f3a346dc8090c567bfcd328deff551c1" +checksum = "d4e88ec9d57c9b685d02f58bfee7be62d72610430ddcedb82a08e5d9925dbfb6" dependencies = [ "arrow", "datafusion-common", @@ -1458,9 +1458,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "159228c3280d342658466bb556dc24de30047fe1d7e559dc5d16ccc5324166f9" +checksum = "8307bb93519b1a91913723a1130cfafeee3f72200d870d88e91a6fc5470ede5c" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1468,9 +1468,9 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5427e5da5edca4d21ea1c7f50e1c9421775fe33d7d5726e5641a833566e7578" +checksum = "2e367e6a71051d0ebdd29b2f85d12059b38b1d1f172c6906e80016da662226bd" dependencies = [ "datafusion-doc", "quote", @@ -1479,9 +1479,9 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89099eefcd5b223ec685c36a41d35c69239236310d71d339f2af0fa4383f3f46" +checksum = "e929015451a67f77d9d8b727b2bf3a40c4445fdef6cdc53281d7d97c76888ace" dependencies = [ "arrow", "chrono", @@ -1499,9 +1499,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f222df5195d605d79098ef37bdd5323bff0131c9d877a24da6ec98dfca9fe36" +checksum = "4b1e68aba7a4b350401cfdf25a3d6f989ad898a7410164afe9ca52080244cb59" dependencies = [ "ahash", "arrow", @@ -1523,9 +1523,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40838625d63d9c12549d81979db3dd675d159055eb9135009ba272ab0e8d0f64" +checksum = "ea22315f33cf2e0adc104e8ec42e285f6ed93998d565c65e82fec6a9ee9f9db4" dependencies = [ "arrow", "datafusion-common", @@ -1538,9 +1538,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eacbcc4cfd502558184ed58fa3c72e775ec65bf077eef5fd2b3453db676f893c" +checksum = "b04b45ea8ad3ac2d78f2ea2a76053e06591c9629c7a603eda16c10649ecf4362" dependencies = [ "ahash", "arrow", @@ -1555,9 +1555,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d501d0e1d0910f015677121601ac177ec59272ef5c9324d1147b394988f40941" +checksum = "7cb13397809a425918f608dfe8653f332015a3e330004ab191b4404187238b95" dependencies = [ "arrow", "datafusion-common", @@ -1574,9 +1574,9 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "463c88ad6f1ecab1810f4c9f046898bee035b370137eb79b2b2db925e270631d" +checksum = "5edc023675791af9d5fb4cc4c24abf5f7bd3bd4dcf9e5bd90ea1eff6976dcc79" dependencies = [ "ahash", "arrow", @@ -1606,9 +1606,9 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677ee4448a010ed5faeff8d73ff78972c2ace59eff3cd7bd15833a1dafa00492" +checksum = "6a387aaef949dc16bb6abc81bd1af850ec7449183aef011214f9724957495738" dependencies = [ "arrow", "chrono", @@ -1629,14 +1629,14 @@ dependencies = [ "datafusion-proto-common", "object_store", "prost", - "rand 0.9.2", + "rand 0.9.4", ] [[package]] name = "datafusion-proto-common" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "965eca01edc8259edbbd95883a00b6d81e329fd44a019cfac3a03b026a83eade" +checksum = "16e614c7c53a9c304c6a850b821010bb492e57300311835f1180613f9d2c63d9" dependencies = [ "arrow", "datafusion-common", @@ -1645,9 +1645,9 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2857618a0ecbd8cd0cf29826889edd3a25774ec26b2995fc3862095c95d88fc6" +checksum = "ac8c76860e355616555081cab5968cec1af7a80701ff374510860bcd567e365a" dependencies = [ "arrow", "datafusion-common", @@ -1705,9 +1705,9 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef8637e35022c5c775003b3ab1debc6b4a8f0eb41b069bdd5475dd3aa93f6eba" +checksum = "5412111aa48e2424ba926112e192f7a6b7e4ccb450145d25ce5ede9f19dc491e" dependencies = [ "async-trait", "datafusion-common", @@ -1719,9 +1719,9 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "53.0.0" +version = "53.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12d9e9f16a1692a11c94bcc418191fa15fd2b4d72a0c1a0c607db93c0b84dd81" +checksum = "fa0d133ddf8b9b3b872acac900157f783e7b879fe9a6bccf389abebbfac45ec1" dependencies = [ "arrow", "bigdecimal", @@ -2083,6 +2083,12 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -2522,12 +2528,11 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.44" +version = "0.1.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" +checksum = "2d1eacfa31c33ec25e873c136ba5669f00f9866d0688bea7be4d3f7e43067df6" dependencies = [ "cc", - "libc", ] [[package]] @@ -2590,9 +2595,9 @@ checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "mimalloc" -version = "0.1.48" +version = "0.1.50" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1ee66a4b64c74f4ef288bcbb9192ad9c3feaad75193129ac8509af543894fd8" +checksum = "b3627c4272df786b9260cabaa46aec1d59c93ede723d4c3ef646c503816b0640" dependencies = [ "libmimalloc-sys", ] @@ -3130,7 +3135,7 @@ dependencies = [ "bytes", "getrandom 0.3.4", "lru-slab", - "rand 0.9.2", + "rand 0.9.4", "ring", "rustc-hash", "rustls", @@ -3179,9 +3184,9 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.9.2" +version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ "rand_chacha", "rand_core 0.9.5", @@ -3428,9 +3433,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.10" +version = "0.103.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df33b2b81ac578cabaf06b89b0631153a3f416b0a886e8a7a1707fb51abbd1ef" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" dependencies = [ "ring", "rustls-pki-types", @@ -3925,9 +3930,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.50.0" +version = "1.52.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "libc", @@ -3940,9 +3945,9 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.6.1" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", @@ -4208,9 +4213,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.23.0" +version = "1.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9" +checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" dependencies = [ "getrandom 0.4.2", "js-sys", diff --git a/Cargo.toml b/Cargo.toml index d0e87a9a4..077bc093f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ members = ["crates/core", "crates/util", "examples/datafusion-ffi-example"] resolver = "3" [workspace.dependencies] -tokio = { version = "1.50" } +tokio = { version = "1.52" } pyo3 = { version = "0.28" } pyo3-async-runtimes = { version = "0.28" } pyo3-log = "0.13.3" From baef8f00e1de6a086fa1814c98e65d216fb16277 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 15 May 2026 09:00:21 -0400 Subject: [PATCH 34/83] Add support for logical and physical codecs (#1541) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: unify logical + physical proto codec stack via SessionContext Introduces a single composable codec layer that every serializer reads from the session, replacing the hardcoded `DefaultLogicalExtensionCodec` / `DefaultPhysicalExtensionCodec` calls scattered across PyLogicalPlan, PyExecutionPlan, and the Rust-wrapped Python provider plumbing. Key changes: * New `PythonLogicalCodec` and `PythonPhysicalCodec` (crates/core/src/codec.rs) wrap any inner `LogicalExtensionCodec` / `PhysicalExtensionCodec`. Both share a `DFPYUDF1` magic-prefix path for in-band cloudpickle encoding of Python scalar UDFs, so an `ExecutionPlan` / `PhysicalExpr` referencing a Python `ScalarUDF` round-trips through either layer. Magic-prefix registry table (DFPYUDF1 in use; DFPYUDA1 / DFPYUDW1 / DFPYPE1 reserved) documented in the module header. * `PySessionContext` stores `Arc` and `Arc` directly. FFI wrappers are built on demand via `ffi_logical_codec()` / `ffi_physical_codec()` for capsule export and downstream `RustWrappedPy*` consumers. Adds `__datafusion_physical_extension_codec__` getter + `with_physical_extension_codec` setter (symmetric with the logical pair). * `PyLogicalPlan.to_proto` / `from_proto` renamed to `to_bytes` / `from_bytes`, now reading the session's logical codec. `to_proto` / `from_proto` survive as deprecated thin wrappers emitting `DeprecationWarning`. * `PyExecutionPlan` gains the same `to_bytes` / `from_bytes` rename + deprecated aliases, plus `__datafusion_execution_plan__` capsule getter and `from_pycapsule` (ported from poc_ffi_query_planner). * New `PyPhysicalExpr` class with `to_bytes` / `from_bytes` / `from_pycapsule` / `__datafusion_physical_expr__`. `from_bytes` takes an input pyarrow Schema for column-reference resolution. * `datafusion-python-util` gains `from_pycapsule!` / `try_from_pycapsule!` macros + `physical_codec_from_pycapsule`, `task_context_from_pycapsule`, `create_physical_extension_capsule` (ported from poc_ffi_query_planner). * `PythonFunctionScalarUDF` exposes `func()`, `input_fields()`, `return_field()`, `volatility()`, `from_parts()` accessors needed by the codec. Python wrapper updates: `LogicalPlan` / `ExecutionPlan` add `to_bytes` / `from_bytes` + deprecate `to_proto` / `from_proto`; `ExecutionPlan` adds capsule getter + `from_pycapsule`; new `PhysicalExpr` wrapper class exported from the top-level package; `SessionContext` exposes the physical codec capsule + setter. Test coverage in python/tests/test_plans.py: round-trip via new API, deprecation warnings on old API, capsule protocol getters, session-routed codec on both layers. `PyLogicalPlan` PyCapsule protocol is intentionally not added — `datafusion-ffi` does not expose `FFI_LogicalPlan`, so there is no stable cross-crate shape to publish. Round-tripping a `LogicalPlan` goes through `to_bytes` / `from_bytes` only. Co-Authored-By: Claude Opus 4.7 (1M context) * test: FFI-example integration tests for codec + plan capsule APIs Adds four downstream-crate fixtures in `datafusion-ffi-example` so the new PR1 surface can be tested with the same FFI-handoff pattern used for table providers, UDFs, etc. Existing tests prove the API exists; these tests prove it composes with code that lives in another crate. New Rust types in `examples/datafusion-ffi-example/src/`: * `MyLogicalExtensionCodec` — delegates to `DefaultLogicalExtensionCodec` and bumps atomic counters on the UDF encode/decode entry points. Exported via `__datafusion_logical_extension_codec__`. Installed onto a session with `ctx.with_logical_extension_codec(my_codec)`. * `MyPhysicalExtensionCodec` — mirror for `PhysicalExtensionCodec`. * `MyExecutionPlan` — wraps a one-column `EmptyExec`, exposes `__datafusion_execution_plan__`. Lets the receiver consume an `ExecutionPlan` capsule that did not originate in datafusion-python. * `MyPhysicalExpr` — wraps `Literal(Int32(42))`, exposes `__datafusion_physical_expr__`. Same FFI handoff for physical expressions. New tests: * `_test_logical_extension_codec.py` — codec installs cleanly, the session re-exports its capsule, and `try_encode_udf` fires on the user codec when serializing a plan that references a `ScalarUDF`. The decode counterpart is a round-trip check rather than a counter assertion: when the UDF is in the receiver's function registry, `parse_expr` resolves by name before consulting the codec. * `_test_physical_extension_codec.py` — symmetric. * `_test_execution_plan.py` — parametrized over typed-class vs raw-capsule input; verifies `ExecutionPlan.from_pycapsule` consumes the downstream capsule. * `_test_physical_expr.py` — same for `PhysicalExpr.from_pycapsule`. API changes forced by the new tests: * `PyLogicalPlan.to_bytes`, `PyExecutionPlan.to_bytes`, `PyPhysicalExpr.to_bytes` now accept an optional `ctx` parameter. When supplied, encoding routes through the session's installed codec instead of a fresh default. `ctx=None` preserves the previous default-codec behavior used by the deprecated `to_proto` shims. * The util `from_pycapsule!` / `try_from_pycapsule!` macros now validate the capsule name via `pointer_checked(Some(c"..."))` rather than `pointer_checked(None)`. The latter rejects named capsules outright with CPython's "incorrect name" error. * `SessionContext.with_logical_extension_codec` and `with_physical_extension_codec` now wrap the returned internal context in `SessionContext` so the result has the full Python surface. The pre-existing logical setter was returning a raw internal object that lacked `sql()` and friends. `examples/datafusion-ffi-example/Cargo.toml` gains `datafusion` and `datafusion-proto` workspace dependencies for the new Rust impls. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: tighten PR1 scope to codec plumbing only Review feedback pass. PR1 is now strictly the composable codec layer + session routing + class-method serialization API. Anything that touches actual Python UDF inline encoding or Python expression wrapping moves to PR2 alongside the pickle work. Dropped: * `encode_python_scalar_udf` / `decode_python_scalar_udf` helpers from `crates/core/src/codec.rs`, along with cloudpickle and pyarrow imports. The wrapper codecs now delegate every method to `inner`. `DFPYUDF1` magic constant is kept (marked `dead_code` for now) as a reservation so PR2 has a single definition site. * `udf.rs` reverted to pre-PR1 shape. The codec no longer needs `func()` / `input_fields()` / `volatility()` / `from_parts()` accessors. Re-added by PR2 when scalar-UDF inlining lands. * `PyPhysicalExpr` class + Python wrapper + `__init__` export + `MyPhysicalExpr` FFI fixture + `_test_physical_expr.py`. No consumer in PR1 or PR2 plan documents; symmetry with `PyExecutionPlan` is not enough to justify the surface area. * Rust-side `PyLogicalPlan::to_proto` / `from_proto` and `PyExecutionPlan::to_proto` / `from_proto` deprecated wrappers. The deprecation lives entirely in the Python wrapper layer, which emits `DeprecationWarning` and forwards to `to_bytes` / `from_bytes`. Less Rust duplication. * `PythonLogicalCodec::with_default_inner` / `PythonPhysicalCodec::with_default_inner` — redundant with `impl Default`. Logic moved into `Default::default`. * `PySessionContext::default_logical_codec` / `default_physical_codec` helpers. Inlined as `Arc::new(PythonLogicalCodec::default())` at the three call sites. Tests (root: 1076, FFI example: 36) all green after the cuts. Co-Authored-By: Claude Opus 4.7 (1M context) * remove unuseful code comments * docs: rewrite codec module comments around purpose, not PR sequence The previous doc-block framed PythonLogicalCodec / PythonPhysicalCodec in terms of "PR1 delegates, PR2 will add encoding" — useful for review, useless for someone reading the code later. Reframed in terms of what the codecs exist to do: encode Python-side plan references (pure-Python UDFs, etc.) into the proto wire format so plans can cross process boundaries without the receiver having to pre-register every callable. The wrappers sit at the top of the session's codec stack and delegate non-Python encoding to a composable inner codec. Magic-prefix registry table loses the "reserved" column. Doc still notes that the in-module impls currently delegate and that encoder/decoder hooks land alongside the corresponding Python-side serialization work. Co-Authored-By: Claude Opus 4.7 (1M context) * feat(codec): forward every LogicalExtensionCodec / PhysicalExtensionCodec method to inner PythonLogicalCodec previously only overrode the four required methods on the trait plus the scalar UDF pair, so the default trait impls (returning "LogicalExtensionCodec is not provided") shadowed any downstream FFI codec for file formats, aggregate UDFs, and window UDFs. A user installing their own codec via `SessionContext.with_logical_extension_codec(...)` would silently lose access to its `try_*_file_format`, `try_*_udaf`, `try_*_udwf` implementations. Forward every trait method to `inner` so the user-installed codec is fully reachable. Same change on the physical side, including `try_*_expr`, `try_*_udaf`, `try_*_udwf` — the corresponding Python-aware paths can layer on later by intercepting before delegation. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: tighten codec dispatch test docstrings The previous docstrings claimed the tests verify "PythonLogicalCodec delegates non-Python UDFs to the inner codec." That's forward-looking — the codecs currently delegate every UDF unconditionally, so the test would behave identically for Python and non-Python UDFs. Rewrite to describe what the test actually proves: the dispatch chain `PyLogicalPlan.to_bytes -> session.logical_codec -> PythonLogicalCodec -> FFI -> user impl` (and the physical mirror) forwards correctly, observable via the user codec's atomic counter incrementing after one encode pass. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor(ffi-example): MyExecutionPlan emits real data via MemorySourceConfig Was a one-column `EmptyExec` stub useful only as a capsule-handoff target. Promoted to a minimal reference impl that a downstream Rust crate can copy when exposing a custom `ExecutionPlan` to datafusion-python: configurable `num_rows`, produces a single batch of sequential `Int32` values under column `value`, wrapped in `DataSourceExec` via `MemorySourceConfig::try_new_exec`. Header comment explains the typical use case (remote backend, streaming source, synthetic data generator) and the `__datafusion_execution_plan__` capsule shape downstream crates should follow. Test asserts the schema-bearing plan survives the FFI hop: a `DataSourceExec` arrives with the expected partitioning and no children. Schema details are not surfaced through the FFI display path (only the wrapping `ForeignExecutionPlan` name + inner plan name appear), so the test does not assert the column name. `to_bytes` round-trip of an FFI-imported plan is not exercised: encoding requires a physical codec that knows how to serialize `ForeignExecutionPlan`, which the default codec does not. A downstream user round-tripping such a plan must install their own codec via `with_physical_extension_codec`. Documented in the test file rather than asserted on. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: drop dormant ExecutionPlan PyCapsule round-trip `PyExecutionPlan::from_pycapsule` and the matching `__datafusion_execution_plan__` exporter have no consumer in this repo, on the POC `poc_ffi_query_planner` branch, or on any sibling branch (`testing/datafusion-distributed`, `testing/ffi-library-marker`, `tmp/ffi-with-codecs`). The pair was wired up speculatively for FFI plan handoff that no Python code path actually performs today. Drop the whole capsule round-trip for `ExecutionPlan`: * Rust `PyExecutionPlan::from_pycapsule` and `__datafusion_execution_plan__`. * Python `ExecutionPlan.from_pycapsule` and `__datafusion_execution_plan__` wrappers. * `MyExecutionPlan` FFI fixture + `_test_execution_plan.py` + lib.rs registration. Was solely a test fixture for the dropped path. * `test_execution_plan_pycapsule_protocol` in `python/tests/test_plans.py`. `PyExecutionPlan.to_bytes` / `from_bytes` survive — they encode through the session's physical codec and have real coverage. Capsule round-trip can be re-added when a concrete consumer (distributed worker, bridge library) lands. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: PyExpr.to_bytes / from_bytes via session logical codec Mirrors PyLogicalPlan / PyExecutionPlan: encode through the session's installed `LogicalExtensionCodec` (or a default-inner `PythonLogicalCodec` when no `ctx` is supplied), decode against the session's function registry + codec via `parse_expr`. Rust side calls `datafusion_proto::logical_plan::to_proto::serialize_expr` and `from_proto::parse_expr`. Python wrapper threads an optional `SessionContext` through. Tests cover the session-routed roundtrip and the no-ctx default-codec encode path. Adds a third consumer of `session.logical_codec()` alongside `PyLogicalPlan` and the codec dispatch tests in the FFI example, broadening coverage of the codec stack. This is the last piece of the PR1 codec surface — follow-up pickle work (`Expr.__reduce__`, worker-scoped context, multiprocessing) can build on this without bundling the byte-level serialization API. Co-Authored-By: Claude Opus 4.7 (1M context) * test(ffi-example): assert codec roundtrip restores plan output PR review feedback: weak `is not None` checks let regressions slip past. Mirror python/tests/test_plans.py — logical compares `df.collect() == round_trip.collect()`; physical compares `str(original) == str(restored)`. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- Cargo.lock | 3 + crates/core/src/codec.rs | 286 ++++++++++++++++++ crates/core/src/context.rs | 135 +++++++-- crates/core/src/expr.rs | 55 ++++ crates/core/src/lib.rs | 1 + crates/core/src/physical_plan.rs | 33 +- crates/core/src/sql/logical.rs | 27 +- crates/util/Cargo.toml | 1 + crates/util/src/lib.rs | 114 +++++++ examples/datafusion-ffi-example/Cargo.toml | 2 + .../tests/_test_logical_extension_codec.py | 82 +++++ .../tests/_test_physical_extension_codec.py | 78 +++++ examples/datafusion-ffi-example/src/lib.rs | 6 + .../src/logical_extension_codec.rs | 153 ++++++++++ .../src/physical_extension_codec.rs | 119 ++++++++ python/datafusion/context.py | 20 +- python/datafusion/expr.py | 20 ++ python/datafusion/plan.py | 84 ++++- python/tests/test_expr.py | 26 ++ python/tests/test_plans.py | 60 +++- 20 files changed, 1235 insertions(+), 70 deletions(-) create mode 100644 crates/core/src/codec.rs create mode 100644 examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py create mode 100644 examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py create mode 100644 examples/datafusion-ffi-example/src/logical_extension_codec.rs create mode 100644 examples/datafusion-ffi-example/src/physical_extension_codec.rs diff --git a/Cargo.lock b/Cargo.lock index 70f09ec46..1d148b0e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1318,12 +1318,14 @@ dependencies = [ "arrow-array", "arrow-schema", "async-trait", + "datafusion", "datafusion-catalog", "datafusion-common", "datafusion-expr", "datafusion-ffi", "datafusion-functions-aggregate", "datafusion-functions-window", + "datafusion-proto", "datafusion-python-util", "pyo3", "pyo3-build-config", @@ -1698,6 +1700,7 @@ dependencies = [ "arrow", "datafusion", "datafusion-ffi", + "datafusion-proto", "prost", "pyo3", "tokio", diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs new file mode 100644 index 000000000..088532df2 --- /dev/null +++ b/crates/core/src/codec.rs @@ -0,0 +1,286 @@ +// 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. + +//! Python-aware extension codecs. +//! +//! Datafusion-python plans can carry references to Python-defined +//! objects that the upstream protobuf codecs do not know how to +//! serialize: pure-Python scalar / aggregate / window UDFs, Python +//! query-planning extensions, and so on. Their state lives inside +//! `Py` callables and closures rather than being recoverable +//! from a name in the receiver's function registry. To ship a plan +//! across a process boundary (pickle, `multiprocessing`, Ray actor, +//! `datafusion-distributed`, etc.) those payloads have to be encoded +//! into the proto wire format itself. +//! +//! [`PythonLogicalCodec`] is the [`LogicalExtensionCodec`] that +//! datafusion-python parks on every `SessionContext`. It wraps a +//! user-supplied (or default) inner codec and adds Python-aware +//! in-band encoding on top: when the encoder sees a Python-defined +//! UDF, the codec cloudpickles the callable + signature into the +//! `fun_definition` proto field; when the decoder sees a payload it +//! produced, it reconstructs the UDF from the bytes alone — no +//! pre-registration on the receiver. UDFs the codec does not +//! recognise are delegated to `inner`, which is typically +//! `DefaultLogicalExtensionCodec` but may be a downstream-supplied +//! FFI codec installed via +//! `SessionContext.with_logical_extension_codec(...)`. +//! +//! [`PythonPhysicalCodec`] is the symmetric wrapper around +//! [`PhysicalExtensionCodec`]. Logical and physical layers each have +//! a `try_encode_udf` / `try_decode_udf` pair, so a `ScalarUDF` +//! referenced inside a `LogicalPlan`, an `ExecutionPlan`, or a +//! `PhysicalExpr` must encode identically through either layer for +//! plans to survive a serialization round-trip. Both codecs share +//! the same payload framing for that reason. +//! +//! Payloads emitted by these codecs are tagged with an 8-byte magic +//! prefix so the decoder can distinguish them from arbitrary bytes +//! (empty `fun_definition` from the default codec, user FFI payloads +//! that picked a non-colliding prefix). Dispatch precedence on +//! decode: **Python-inline payload (magic prefix match) → `inner` +//! codec → caller's `FunctionRegistry` fallback.** +//! +//! ## Wire-format magic prefix registry +//! +//! | Layer + kind | Magic prefix | +//! | ----------------------------- | ------------ | +//! | `PythonLogicalCodec` scalar | `DFPYUDF1` | +//! | `PythonLogicalCodec` agg | `DFPYUDA1` | +//! | `PythonLogicalCodec` window | `DFPYUDW1` | +//! | `PythonPhysicalCodec` scalar | `DFPYUDF1` | +//! | `PythonPhysicalCodec` agg | `DFPYUDA1` | +//! | `PythonPhysicalCodec` window | `DFPYUDW1` | +//! | `PythonPhysicalCodec` expr | `DFPYPE1` | +//! | User FFI extension codec | user-chosen | +//! | Default codec | (none) | +//! +//! Downstream FFI codecs should pick non-colliding prefixes (use a +//! `DF` namespace plus a crate-specific suffix). The codec +//! implementations in this module currently delegate every method to +//! `inner`; the encoder/decoder hooks for each kind are added as the +//! corresponding Python-side type becomes serializable. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion::common::{Result, TableReference}; +use datafusion::datasource::TableProvider; +use datafusion::datasource::file_format::FileFormatFactory; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::{AggregateUDF, Extension, LogicalPlan, ScalarUDF, WindowUDF}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_proto::physical_plan::{DefaultPhysicalExtensionCodec, PhysicalExtensionCodec}; + +/// Wire-format prefix that tags a `fun_definition` payload as an +/// inlined Python scalar UDF (cloudpickled tuple of name, callable, +/// input schema, return field, volatility). Defined once here so +/// the encoder and decoder cannot drift. +#[allow(dead_code)] +pub(crate) const PY_SCALAR_UDF_MAGIC: &[u8] = b"DFPYUDF1"; + +/// `LogicalExtensionCodec` parked on every `SessionContext`. Holds +/// the Python-aware encoding hooks for logical-layer types +/// (`LogicalPlan`, `Expr`) and delegates everything it does not +/// handle to the composable `inner` codec — typically +/// `DefaultLogicalExtensionCodec`, or a downstream FFI codec +/// installed via `SessionContext.with_logical_extension_codec(...)`. +/// +/// Sitting at the top of the session's logical codec stack means +/// every serializer that reads `session.logical_codec()` automatically +/// picks up Python-aware encoding for free. +#[derive(Debug)] +pub struct PythonLogicalCodec { + inner: Arc, +} + +impl PythonLogicalCodec { + pub fn new(inner: Arc) -> Self { + Self { inner } + } + + pub fn inner(&self) -> &Arc { + &self.inner + } +} + +impl Default for PythonLogicalCodec { + fn default() -> Self { + Self::new(Arc::new(DefaultLogicalExtensionCodec {})) + } +} + +impl LogicalExtensionCodec for PythonLogicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[LogicalPlan], + ctx: &TaskContext, + ) -> Result { + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: SchemaRef, + ctx: &TaskContext, + ) -> Result> { + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode_table_provider(table_ref, node, buf) + } + + fn try_decode_file_format( + &self, + buf: &[u8], + ctx: &TaskContext, + ) -> Result> { + self.inner.try_decode_file_format(buf, ctx) + } + + fn try_encode_file_format( + &self, + buf: &mut Vec, + node: Arc, + ) -> Result<()> { + self.inner.try_encode_file_format(buf, node) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + self.inner.try_encode_udf(node, buf) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + self.inner.try_decode_udf(name, buf) + } + + fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { + self.inner.try_encode_udaf(node, buf) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + self.inner.try_decode_udaf(name, buf) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + self.inner.try_encode_udwf(node, buf) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + self.inner.try_decode_udwf(name, buf) + } +} + +/// `PhysicalExtensionCodec` mirror of [`PythonLogicalCodec`] parked +/// on the same `SessionContext`. Carries the Python-aware encoding +/// hooks for physical-layer types (`ExecutionPlan`, `PhysicalExpr`) +/// and delegates the rest to `inner`. +/// +/// The `PhysicalExtensionCodec` trait has its own `try_encode_udf` +/// / `try_decode_udf` pair distinct from the logical one, so a +/// `ScalarUDF` referenced inside a physical plan needs Python-aware +/// encoding on this layer too — otherwise a plan with a Python UDF +/// would round-trip at the logical level but break at the physical +/// level. Both layers reuse the shared payload framing +/// ([`PY_SCALAR_UDF_MAGIC`] et al.) so the wire format is identical. +#[derive(Debug)] +pub struct PythonPhysicalCodec { + inner: Arc, +} + +impl PythonPhysicalCodec { + pub fn new(inner: Arc) -> Self { + Self { inner } + } + + pub fn inner(&self) -> &Arc { + &self.inner + } +} + +impl Default for PythonPhysicalCodec { + fn default() -> Self { + Self::new(Arc::new(DefaultPhysicalExtensionCodec {})) + } +} + +impl PhysicalExtensionCodec for PythonPhysicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + ) -> Result> { + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + self.inner.try_encode_udf(node, buf) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + self.inner.try_decode_udf(name, buf) + } + + fn try_encode_expr(&self, node: &Arc, buf: &mut Vec) -> Result<()> { + self.inner.try_encode_expr(node, buf) + } + + fn try_decode_expr( + &self, + buf: &[u8], + inputs: &[Arc], + ) -> Result> { + self.inner.try_decode_expr(buf, inputs) + } + + fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { + self.inner.try_encode_udaf(node, buf) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + self.inner.try_decode_udaf(name, buf) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + self.inner.try_encode_udwf(node, buf) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + self.inner.try_decode_udwf(name, buf) + } +} diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index e46d359d6..96de01889 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -52,11 +52,14 @@ use datafusion_ffi::catalog_provider_list::FFI_CatalogProviderList; use datafusion_ffi::config::extension_options::FFI_ExtensionOptions; use datafusion_ffi::execution::FFI_TaskContextProvider; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; -use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; +use datafusion_proto::logical_plan::LogicalExtensionCodec; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ - create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx, - get_tokio_runtime, spawn_future, wait_for_future, + create_logical_extension_capsule, create_physical_extension_capsule, + ffi_logical_codec_from_pycapsule, get_global_ctx, get_tokio_runtime, + physical_codec_from_pycapsule, spawn_future, wait_for_future, }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; @@ -69,6 +72,7 @@ use uuid::Uuid; use crate::catalog::{ PyCatalog, PyCatalogList, RustWrappedPyCatalogProvider, RustWrappedPyCatalogProviderList, }; +use crate::codec::{PythonLogicalCodec, PythonPhysicalCodec}; use crate::common::data_type::PyScalarValue; use crate::common::df_schema::PyDFSchema; use crate::dataframe::PyDataFrame; @@ -365,7 +369,8 @@ impl PySQLOptions { #[derive(Clone)] pub struct PySessionContext { pub ctx: Arc, - logical_codec: Arc, + logical_codec: Arc, + physical_codec: Arc, } #[pymethods] @@ -393,14 +398,18 @@ impl PySessionContext { .with_default_features() .build(); let ctx = Arc::new(SessionContext::new_with_state(session_state)); - let logical_codec = Self::default_logical_codec(&ctx); - Ok(PySessionContext { ctx, logical_codec }) + Ok(PySessionContext { + ctx, + logical_codec: Arc::new(PythonLogicalCodec::default()), + physical_codec: Arc::new(PythonPhysicalCodec::default()), + }) } pub fn enable_url_table(&self) -> PyResult { Ok(PySessionContext { ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()), logical_codec: Arc::clone(&self.logical_codec), + physical_codec: Arc::clone(&self.physical_codec), }) } @@ -408,8 +417,11 @@ impl PySessionContext { #[pyo3(signature = ())] pub fn global_ctx() -> PyResult { let ctx = get_global_ctx().clone(); - let logical_codec = Self::default_logical_codec(&ctx); - Ok(Self { ctx, logical_codec }) + Ok(Self { + ctx, + logical_codec: Arc::new(PythonLogicalCodec::default()), + physical_codec: Arc::new(PythonPhysicalCodec::default()), + }) } /// Register an object store with the given name @@ -714,7 +726,8 @@ impl PySessionContext { ) -> PyDataFusionResult<()> { if factory.hasattr("__datafusion_table_provider_factory__")? { let py = factory.py(); - let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?; + let ffi = self.ffi_logical_codec(); + let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; factory = factory .getattr("__datafusion_table_provider_factory__")? .call1((codec_capsule,))?; @@ -730,7 +743,7 @@ impl PySessionContext { } else { Arc::new(RustWrappedPyTableProviderFactory::new( factory.into(), - self.logical_codec.clone(), + self.ffi_logical_codec(), )) }; @@ -748,7 +761,8 @@ impl PySessionContext { ) -> PyDataFusionResult<()> { if provider.hasattr("__datafusion_catalog_provider_list__")? { let py = provider.py(); - let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?; + let ffi = self.ffi_logical_codec(); + let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; provider = provider .getattr("__datafusion_catalog_provider_list__")? .call1((codec_capsule,))?; @@ -766,7 +780,7 @@ impl PySessionContext { Ok(py_catalog_list) => py_catalog_list.catalog_list, Err(_) => Arc::new(RustWrappedPyCatalogProviderList::new( provider.into(), - Arc::clone(&self.logical_codec), + self.ffi_logical_codec(), )) as Arc, } }; @@ -783,7 +797,8 @@ impl PySessionContext { ) -> PyDataFusionResult<()> { if provider.hasattr("__datafusion_catalog_provider__")? { let py = provider.py(); - let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?; + let ffi = self.ffi_logical_codec(); + let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; provider = provider .getattr("__datafusion_catalog_provider__")? .call1((codec_capsule,))?; @@ -801,7 +816,7 @@ impl PySessionContext { Ok(py_catalog) => py_catalog.catalog, Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( provider.into(), - Arc::clone(&self.logical_codec), + self.ffi_logical_codec(), )) as Arc, } }; @@ -1061,10 +1076,9 @@ impl PySessionContext { .downcast_ref::() { Some(wrapped_schema) => Ok(wrapped_schema.catalog_provider.clone_ref(py)), - None => Ok( - PyCatalog::new_from_parts(catalog, Arc::clone(&self.logical_codec)) - .into_py_any(py)?, - ), + None => { + Ok(PyCatalog::new_from_parts(catalog, self.ffi_logical_codec()).into_py_any(py)?) + } } } @@ -1353,20 +1367,44 @@ impl PySessionContext { &self, py: Python<'py>, ) -> PyResult> { - create_logical_extension_capsule(py, self.logical_codec.as_ref()) + let ffi = self.ffi_logical_codec(); + create_logical_extension_capsule(py, ffi.as_ref()) } pub fn with_logical_extension_codec<'py>( &self, codec: Bound<'py, PyAny>, ) -> PyDataFusionResult { - let logical_codec = Arc::new(ffi_logical_codec_from_pycapsule(codec)?); + let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?; + let inner: Arc = (&inner_ffi).into(); + let logical_codec = Arc::new(PythonLogicalCodec::new(inner)); + + Ok(Self { + ctx: Arc::clone(&self.ctx), + logical_codec, + physical_codec: Arc::clone(&self.physical_codec), + }) + } - Ok({ - Self { - ctx: Arc::clone(&self.ctx), - logical_codec, - } + pub fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let ffi = self.ffi_physical_codec(); + create_physical_extension_capsule(py, ffi.as_ref()) + } + + pub fn with_physical_extension_codec<'py>( + &self, + codec: Bound<'py, PyAny>, + ) -> PyDataFusionResult { + let inner = physical_codec_from_pycapsule(&codec)?; + let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); + + Ok(Self { + ctx: Arc::clone(&self.ctx), + logical_codec: Arc::clone(&self.logical_codec), + physical_codec, }) } } @@ -1416,12 +1454,42 @@ impl PySessionContext { Ok(()) } - fn default_logical_codec(ctx: &Arc) -> Arc { - let codec = Arc::new(DefaultLogicalExtensionCodec {}); + /// Session-scoped logical codec. Sibling modules read this when they + /// need to serialize/deserialize logical-layer types (LogicalPlan, + /// Expr) against the user-installed (or default) codec stack. + pub(crate) fn logical_codec(&self) -> &Arc { + &self.logical_codec + } + + /// Session-scoped physical codec. Sibling modules read this for + /// ExecutionPlan / PhysicalExpr serialization. + pub(crate) fn physical_codec(&self) -> &Arc { + &self.physical_codec + } + + /// Build an FFI-wrapped clone of the session's logical codec on demand. + /// Used at every site that exports the codec across an FFI boundary + /// (capsule getters, Rust wrappers for Python-defined providers, etc.). + pub(crate) fn ffi_logical_codec(&self) -> Arc { + let inner: Arc = + Arc::clone(&self.logical_codec) as Arc; let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::clone(ctx) as Arc; + let ctx_provider = Arc::clone(&self.ctx) as Arc; Arc::new(FFI_LogicalExtensionCodec::new( - codec, + inner, + Some(runtime), + &ctx_provider, + )) + } + + /// Build an FFI-wrapped clone of the session's physical codec on demand. + pub(crate) fn ffi_physical_codec(&self) -> Arc { + let inner: Arc = + Arc::clone(&self.physical_codec) as Arc; + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = Arc::clone(&self.ctx) as Arc; + Arc::new(FFI_PhysicalExtensionCodec::new( + inner, Some(runtime), &ctx_provider, )) @@ -1445,9 +1513,10 @@ impl From for SessionContext { impl From for PySessionContext { fn from(ctx: SessionContext) -> PySessionContext { - let ctx = Arc::new(ctx); - let logical_codec = Self::default_logical_codec(&ctx); - - PySessionContext { ctx, logical_codec } + PySessionContext { + ctx: Arc::new(ctx), + logical_codec: Arc::new(PythonLogicalCodec::default()), + physical_codec: Arc::new(PythonPhysicalCodec::default()), + } } } diff --git a/crates/core/src/expr.rs b/crates/core/src/expr.rs index c4f2a12da..2e633baeb 100644 --- a/crates/core/src/expr.rs +++ b/crates/core/src/expr.rs @@ -31,9 +31,13 @@ use datafusion::logical_expr::{ Between, BinaryExpr, Case, Cast, Expr, ExprFuncBuilder, ExprFunctionExt, Like, LogicalPlan, Operator, TryCast, WindowFunctionDefinition, col, lit, lit_with_metadata, }; +use datafusion_proto::logical_plan::{from_proto, to_proto}; +use prost::Message; use pyo3::IntoPyObjectExt; use pyo3::basic::CompareOp; +use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; +use pyo3::types::PyBytes; use window::PyWindowFrame; use self::alias::PyAlias; @@ -43,7 +47,9 @@ use self::bool_expr::{ }; use self::like::{PyILike, PyLike, PySimilarTo}; use self::scalar_variable::PyScalarVariable; +use crate::codec::PythonLogicalCodec; use crate::common::data_type::{DataTypeMap, NullTreatment, PyScalarValue, RexType}; +use crate::context::PySessionContext; use crate::errors::{PyDataFusionResult, py_runtime_err, py_type_err, py_unsupported_variant_err}; use crate::expr::aggregate_expr::PyAggregateFunction; use crate::expr::binary_expr::PyBinaryExpr; @@ -660,6 +666,55 @@ impl PyExpr { .into()), } } + + /// Serialize this `Expr` to protobuf bytes. + /// + /// When `ctx` is supplied, encoding routes through the session's + /// installed `LogicalExtensionCodec` so user FFI codecs see the + /// encode path. Without `ctx` a default-inner Python codec is + /// used; Python scalar UDFs still inline when in-band encoding + /// lands, non-Python UDFs fall through to the default codec. + #[pyo3(signature = (ctx=None))] + pub fn to_bytes<'py>( + &'py self, + py: Python<'py>, + ctx: Option, + ) -> PyDataFusionResult> { + let default_codec; + let codec: &dyn datafusion_proto::logical_plan::LogicalExtensionCodec = match ctx { + Some(ref ctx) => ctx.logical_codec().as_ref(), + None => { + default_codec = PythonLogicalCodec::default(); + &default_codec + } + }; + let proto = to_proto::serialize_expr(&self.expr, codec) + .map_err(|e| PyRuntimeError::new_err(format!("Unable to serialize expr: {e}")))?; + let bytes = proto.encode_to_vec(); + Ok(PyBytes::new(py, &bytes)) + } + + /// Decode an `Expr` from protobuf bytes against the session's + /// function registry and logical codec. + #[staticmethod] + pub fn from_bytes( + ctx: PySessionContext, + proto_msg: Bound<'_, PyBytes>, + ) -> PyDataFusionResult { + let bytes: &[u8] = proto_msg.extract().map_err(Into::::into)?; + let proto_expr = + datafusion_proto::protobuf::LogicalExprNode::decode(bytes).map_err(|e| { + PyRuntimeError::new_err(format!( + "Unable to decode expression from serialized bytes: {e}" + )) + })?; + + let codec = ctx.logical_codec(); + let task_ctx = ctx.ctx.task_ctx(); + let expr = from_proto::parse_expr(&proto_expr, task_ctx.as_ref(), codec.as_ref()) + .map_err(|e| PyRuntimeError::new_err(format!("Unable to decode expr: {e}")))?; + Ok(Self { expr }) + } } #[pyclass( diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 77d69911a..e3551c937 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -28,6 +28,7 @@ use pyo3::prelude::*; #[allow(clippy::borrow_deref_ref)] pub mod catalog; +pub mod codec; pub mod common; #[allow(clippy::borrow_deref_ref)] diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index fac973884..594655a60 100644 --- a/crates/core/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -18,12 +18,13 @@ use std::sync::Arc; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; -use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; +use datafusion_proto::physical_plan::AsExecutionPlan; use prost::Message; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyBytes; +use crate::codec::PythonPhysicalCodec; use crate::context::PySessionContext; use crate::errors::PyDataFusionResult; use crate::metrics::PyMetricsSet; @@ -68,11 +69,26 @@ impl PyExecutionPlan { format!("{}", d.indent(false)) } - pub fn to_proto<'py>(&'py self, py: Python<'py>) -> PyDataFusionResult> { - let codec = DefaultPhysicalExtensionCodec {}; + #[pyo3(signature = (ctx=None))] + pub fn to_bytes<'py>( + &'py self, + py: Python<'py>, + ctx: Option, + ) -> PyDataFusionResult> { + // Route through the session's physical codec when supplied so + // user FFI codecs registered via + // `with_physical_extension_codec` see the encode path. + let default_codec; + let codec: &dyn datafusion_proto::physical_plan::PhysicalExtensionCodec = match ctx { + Some(ref ctx) => ctx.physical_codec().as_ref(), + None => { + default_codec = PythonPhysicalCodec::default(); + &default_codec + } + }; let proto = datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan( self.plan.clone(), - &codec, + codec, )?; let bytes = proto.encode_to_vec(); @@ -80,7 +96,7 @@ impl PyExecutionPlan { } #[staticmethod] - pub fn from_proto( + pub fn from_bytes( ctx: PySessionContext, proto_msg: Bound<'_, PyBytes>, ) -> PyDataFusionResult { @@ -88,12 +104,13 @@ impl PyExecutionPlan { let proto_plan = datafusion_proto::protobuf::PhysicalPlanNode::decode(bytes).map_err(|e| { PyRuntimeError::new_err(format!( - "Unable to decode logical node from serialized bytes: {e}" + "Unable to decode physical node from serialized bytes: {e}" )) })?; - let codec = DefaultPhysicalExtensionCodec {}; - let plan = proto_plan.try_into_physical_plan(ctx.ctx.task_ctx().as_ref(), &codec)?; + let codec = ctx.physical_codec(); + let plan = + proto_plan.try_into_physical_plan(ctx.ctx.task_ctx().as_ref(), codec.as_ref())?; Ok(Self::new(plan)) } diff --git a/crates/core/src/sql/logical.rs b/crates/core/src/sql/logical.rs index 631aa9b09..647c3fa7e 100644 --- a/crates/core/src/sql/logical.rs +++ b/crates/core/src/sql/logical.rs @@ -18,12 +18,13 @@ use std::sync::Arc; use datafusion::logical_expr::{DdlStatement, LogicalPlan, Statement}; -use datafusion_proto::logical_plan::{AsLogicalPlan, DefaultLogicalExtensionCodec}; +use datafusion_proto::logical_plan::AsLogicalPlan; use prost::Message; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyBytes; +use crate::codec::PythonLogicalCodec; use crate::context::PySessionContext; use crate::errors::PyDataFusionResult; use crate::expr::aggregate::PyAggregate; @@ -196,17 +197,29 @@ impl PyLogicalPlan { format!("{}", self.plan.display_graphviz()) } - pub fn to_proto<'py>(&'py self, py: Python<'py>) -> PyDataFusionResult> { - let codec = DefaultLogicalExtensionCodec {}; + #[pyo3(signature = (ctx=None))] + pub fn to_bytes<'py>( + &'py self, + py: Python<'py>, + ctx: Option, + ) -> PyDataFusionResult> { + let default_codec; + let codec: &dyn datafusion_proto::logical_plan::LogicalExtensionCodec = match ctx { + Some(ref ctx) => ctx.logical_codec().as_ref(), + None => { + default_codec = PythonLogicalCodec::default(); + &default_codec + } + }; let proto = - datafusion_proto::protobuf::LogicalPlanNode::try_from_logical_plan(&self.plan, &codec)?; + datafusion_proto::protobuf::LogicalPlanNode::try_from_logical_plan(&self.plan, codec)?; let bytes = proto.encode_to_vec(); Ok(PyBytes::new(py, &bytes)) } #[staticmethod] - pub fn from_proto( + pub fn from_bytes( ctx: PySessionContext, proto_msg: Bound<'_, PyBytes>, ) -> PyDataFusionResult { @@ -218,8 +231,8 @@ impl PyLogicalPlan { )) })?; - let codec = DefaultLogicalExtensionCodec {}; - let plan = proto_plan.try_into_logical_plan(&ctx.ctx.task_ctx(), &codec)?; + let codec = ctx.logical_codec(); + let plan = proto_plan.try_into_logical_plan(&ctx.ctx.task_ctx(), codec.as_ref())?; Ok(Self::new(plan)) } } diff --git a/crates/util/Cargo.toml b/crates/util/Cargo.toml index 00d5946a5..c23667b0f 100644 --- a/crates/util/Cargo.toml +++ b/crates/util/Cargo.toml @@ -30,5 +30,6 @@ tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] } pyo3 = { workspace = true } datafusion = { workspace = true } datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } arrow = { workspace = true } prost = { workspace = true } diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 5b1c89936..72dc9aafc 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -21,10 +21,14 @@ use std::sync::{Arc, OnceLock}; use std::time::Duration; use datafusion::datasource::TableProvider; +use datafusion::execution::TaskContext; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::Volatility; +use datafusion_ffi::execution::FFI_TaskContextProvider; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_ffi::table_provider::FFI_TableProvider; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyType}; @@ -224,3 +228,113 @@ pub fn ffi_logical_codec_from_pycapsule(obj: Bound) -> PyResult( + py: Python<'py>, + codec: &FFI_PhysicalExtensionCodec, +) -> PyResult> { + let name = cr"datafusion_physical_extension_codec".into(); + let codec = codec.clone(); + + PyCapsule::new(py, codec, Some(name)) +} + +/// Define a `(obj) -> PyResult>` extractor that +/// accepts either a raw `PyCapsule` carrying `$ffi_type` or any object +/// exposing `____()` that returns one. +/// +/// Use this when `Arc<$output_type>: From<&$ffi_type>` (infallible +/// conversion). For fallible conversions use [`try_from_pycapsule!`] +/// instead. +#[macro_export] +macro_rules! from_pycapsule { + ($fn_name:ident, $capsule_name:literal, $ffi_type:ty, $output_type:ty) => { + pub fn $fn_name( + obj: &$crate::pyo3::Bound<$crate::pyo3::PyAny>, + ) -> $crate::pyo3::PyResult> { + use $crate::pyo3::prelude::*; + use $crate::pyo3::types::PyCapsule; + + let mut obj = obj.clone(); + if obj.hasattr(concat!("__", $capsule_name, "__"))? { + obj = obj.getattr(concat!("__", $capsule_name, "__"))?.call0()?; + } + let capsule = obj.cast::().map_err(|_| { + $crate::errors::py_datafusion_err(concat!( + "Invalid ", + $capsule_name, + ". Does not contain PyCapsule object." + )) + })?; + $crate::validate_pycapsule(&capsule, $capsule_name)?; + + let expected_name = std::ffi::CString::new($capsule_name) + .expect("capsule name must not contain interior NUL bytes"); + let data: std::ptr::NonNull<$ffi_type> = capsule + .pointer_checked(Some(expected_name.as_c_str()))? + .cast(); + let output_obj = unsafe { data.as_ref() }; + let output_obj: std::sync::Arc<$output_type> = output_obj.into(); + + Ok(output_obj) + } + }; +} + +/// Same shape as [`from_pycapsule!`] but for FFI types whose conversion +/// into `Arc<$output_type>` is fallible (uses `TryFrom`). +#[macro_export] +macro_rules! try_from_pycapsule { + ($fn_name:ident, $capsule_name:literal, $ffi_type:ty, $output_type:ty) => { + pub fn $fn_name( + obj: &$crate::pyo3::Bound<$crate::pyo3::PyAny>, + ) -> $crate::pyo3::PyResult> { + use $crate::pyo3::prelude::*; + use $crate::pyo3::types::PyCapsule; + + let mut obj = obj.clone(); + if obj.hasattr(concat!("__", $capsule_name, "__"))? { + obj = obj.getattr(concat!("__", $capsule_name, "__"))?.call0()?; + } + let capsule = obj.cast::().map_err(|_| { + $crate::errors::py_datafusion_err(concat!( + "Invalid ", + $capsule_name, + ". Does not contain PyCapsule object." + )) + })?; + $crate::validate_pycapsule(&capsule, $capsule_name)?; + + let expected_name = std::ffi::CString::new($capsule_name) + .expect("capsule name must not contain interior NUL bytes"); + let data: std::ptr::NonNull<$ffi_type> = capsule + .pointer_checked(Some(expected_name.as_c_str()))? + .cast(); + let output_obj = unsafe { data.as_ref() }; + let output_obj: std::sync::Arc<$output_type> = output_obj + .try_into() + .map_err($crate::errors::py_datafusion_err)?; + + Ok(output_obj) + } + }; +} + +// Re-export pyo3 so the macros expand inside downstream crates without +// requiring an explicit pyo3 dep at the call site. +#[doc(hidden)] +pub use pyo3; + +from_pycapsule!( + physical_codec_from_pycapsule, + "datafusion_physical_extension_codec", + FFI_PhysicalExtensionCodec, + dyn PhysicalExtensionCodec +); + +try_from_pycapsule!( + task_context_from_pycapsule, + "datafusion_task_context_provider", + FFI_TaskContextProvider, + TaskContext +); diff --git a/examples/datafusion-ffi-example/Cargo.toml b/examples/datafusion-ffi-example/Cargo.toml index 178dce9f9..ffc839d56 100644 --- a/examples/datafusion-ffi-example/Cargo.toml +++ b/examples/datafusion-ffi-example/Cargo.toml @@ -26,12 +26,14 @@ repository.workspace = true publish = false [dependencies] +datafusion = { workspace = true } datafusion-catalog = { workspace = true, default-features = false } datafusion-common = { workspace = true, default-features = false } datafusion-functions-aggregate = { workspace = true } datafusion-functions-window = { workspace = true } datafusion-expr = { workspace = true } datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } arrow = { workspace = true } arrow-array = { workspace = true } diff --git a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py new file mode 100644 index 000000000..cd0c5a61a --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py @@ -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. + +from __future__ import annotations + +from datafusion import LogicalPlan, SessionContext +from datafusion_ffi_example import MyLogicalExtensionCodec + + +def _setup_session_with_codec() -> tuple[SessionContext, MyLogicalExtensionCodec]: + """Build a session with the user-supplied logical extension codec + installed. Tests use a FROM-less query so plan serialization does + not pull in `try_encode_table_provider`, which the default codec + leaves unimplemented.""" + base = SessionContext() + codec = MyLogicalExtensionCodec() + ctx = base.with_logical_extension_codec(codec) + return ctx, codec + + +def test_ffi_logical_codec_install_and_export(): + """Installing a user FFI codec replaces the session's logical + codec; the capsule getter on the session re-exports it.""" + ctx, _codec = _setup_session_with_codec() + capsule = ctx.__datafusion_logical_extension_codec__() + assert capsule is not None + + +def test_ffi_logical_codec_consulted_on_udf_encode(): + """Serializing through ctx.logical_codec() routes try_encode_udf to + the user-installed FFI codec. + + Verifies the dispatch chain + `PyLogicalPlan.to_bytes -> session.logical_codec -> + PythonLogicalCodec -> FFI_LogicalExtensionCodec -> user impl` + is wired correctly. The user codec's atomic counter increments + after a serialization pass, proving every hop forwards. + + Does not test any Python-UDF-specific dispatch — PythonLogicalCodec + currently delegates all UDF encoding to its inner codec + unconditionally. Python-vs-other branching lands when in-band + scalar UDF encoding is added. + """ + ctx, codec = _setup_session_with_codec() + df = ctx.sql("SELECT abs(-1) AS x") + plan = df.logical_plan() + + before = codec.encode_udf_calls() + _ = plan.to_bytes(ctx) + after = codec.encode_udf_calls() + + assert after > before, ( + f"Expected user FFI codec encode_udf to fire, before={before} after={after}" + ) + + +def test_ffi_logical_codec_roundtrip(): + """A plan referencing an FFI-imported UDF round-trips through the + user-supplied logical codec (encode via codec, decode resolves from + registry — `try_decode_udf` is only consulted when the UDF is not + in the registry, which is the codec-inlined case).""" + ctx, _codec = _setup_session_with_codec() + df = ctx.sql("SELECT abs(-1) AS x") + blob = df.logical_plan().to_bytes(ctx) + + restored = LogicalPlan.from_bytes(ctx, blob) + df_round_trip = ctx.create_dataframe_from_logical_plan(restored) + assert df.collect() == df_round_trip.collect() diff --git a/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py new file mode 100644 index 000000000..28eaaf2d9 --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py @@ -0,0 +1,78 @@ +# 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. + +from __future__ import annotations + +import pyarrow as pa +from datafusion import ExecutionPlan, SessionContext +from datafusion_ffi_example import MyPhysicalExtensionCodec + + +def _setup_session_with_codec() -> tuple[SessionContext, MyPhysicalExtensionCodec]: + base = SessionContext() + batch = pa.RecordBatch.from_arrays( + [pa.array([-1, -2, -3])], + names=["a"], + ) + base.register_record_batches("t", [[batch]]) + codec = MyPhysicalExtensionCodec() + ctx = base.with_physical_extension_codec(codec) + return ctx, codec + + +def test_ffi_physical_codec_install_and_export(): + ctx, _codec = _setup_session_with_codec() + capsule = ctx.__datafusion_physical_extension_codec__() + assert capsule is not None + + +def test_ffi_physical_codec_consulted_on_udf_encode(): + """Serializing through ctx.physical_codec() routes try_encode_udf to + the user-installed FFI codec. + + Mirror of the logical-side dispatch test: verifies + `PyExecutionPlan.to_bytes -> session.physical_codec -> + PythonPhysicalCodec -> FFI_PhysicalExtensionCodec -> user impl` + forwards correctly. Does not test Python-UDF-specific dispatch — + PythonPhysicalCodec currently delegates all UDF encoding to its + inner codec unconditionally. + """ + ctx, codec = _setup_session_with_codec() + df = ctx.sql("SELECT abs(a) AS x FROM t") + plan = df.execution_plan() + + before = codec.encode_udf_calls() + _ = plan.to_bytes(ctx) + after = codec.encode_udf_calls() + + assert after > before, ( + f"Expected user FFI codec encode_udf to fire, before={before} after={after}" + ) + + +def test_ffi_physical_codec_roundtrip(): + """A plan referencing an FFI-imported UDF round-trips via the + user-supplied physical codec. On decode, the receiver resolves the + UDF from the function registry; `try_decode_udf` only fires when a + codec inlines the UDF body, which the counting codec does not.""" + ctx, _codec = _setup_session_with_codec() + df = ctx.sql("SELECT abs(a) AS x FROM t") + original = df.execution_plan() + blob = original.to_bytes(ctx) + + restored = ExecutionPlan.from_bytes(ctx, blob) + assert str(original) == str(restored) diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index e708c49cc..3323ac982 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -20,6 +20,8 @@ use pyo3::prelude::*; use crate::aggregate_udf::MySumUDF; use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList}; use crate::config::MyConfig; +use crate::logical_extension_codec::MyLogicalExtensionCodec; +use crate::physical_extension_codec::MyPhysicalExtensionCodec; use crate::scalar_udf::IsNullUDF; use crate::table_function::MyTableFunction; use crate::table_provider::MyTableProvider; @@ -29,6 +31,8 @@ use crate::window_udf::MyRankUDF; pub(crate) mod aggregate_udf; pub(crate) mod catalog_provider; pub(crate) mod config; +pub(crate) mod logical_extension_codec; +pub(crate) mod physical_extension_codec; pub(crate) mod scalar_udf; pub(crate) mod table_function; pub(crate) mod table_provider; @@ -49,5 +53,7 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs new file mode 100644 index 000000000..da9efb297 --- /dev/null +++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs @@ -0,0 +1,153 @@ +// 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. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow::datatypes::SchemaRef; +use datafusion::common::{Result, TableReference}; +use datafusion::datasource::TableProvider; +use datafusion::execution::{TaskContext, TaskContextProvider}; +use datafusion::logical_expr::{Extension, LogicalPlan, ScalarUDF}; +use datafusion::prelude::SessionContext; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_python_util::get_tokio_runtime; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +/// Tracks how often each `try_*_udf` entry point fires. Surface for +/// Python tests to assert the session routed UDF +/// encode/decode through this user-supplied codec rather than the +/// upstream default. +#[derive(Debug, Default)] +pub(crate) struct CallCounters { + pub encode_udf: AtomicUsize, + pub decode_udf: AtomicUsize, +} + +/// Minimal user-supplied `LogicalExtensionCodec` for integration tests. +/// Delegates everything to `DefaultLogicalExtensionCodec` and bumps +/// counters on the UDF entry points so tests can prove the wrapper +/// installed via `SessionContext.with_logical_extension_codec(...)` +/// actually gets consulted. +#[derive(Debug)] +struct CountingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec, + counters: Arc, +} + +impl LogicalExtensionCodec for CountingLogicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[LogicalPlan], + ctx: &TaskContext, + ) -> Result { + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: SchemaRef, + ctx: &TaskContext, + ) -> Result> { + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode_table_provider(table_ref, node, buf) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + self.counters.decode_udf.fetch_add(1, Ordering::SeqCst); + self.inner.try_decode_udf(name, buf) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + self.counters.encode_udf.fetch_add(1, Ordering::SeqCst); + self.inner.try_encode_udf(node, buf) + } +} + +#[pyclass( + from_py_object, + name = "MyLogicalExtensionCodec", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Clone)] +pub(crate) struct MyLogicalExtensionCodec { + counters: Arc, +} + +#[pymethods] +impl MyLogicalExtensionCodec { + #[new] + fn new() -> Self { + Self { + counters: Arc::new(CallCounters::default()), + } + } + + /// Number of `try_encode_udf` invocations observed since + /// construction. + fn encode_udf_calls(&self) -> usize { + self.counters.encode_udf.load(Ordering::SeqCst) + } + + /// Number of `try_decode_udf` invocations observed. + fn decode_udf_calls(&self) -> usize { + self.counters.decode_udf.load(Ordering::SeqCst) + } + + /// Capsule entry point consumed by + /// `datafusion_python_util::ffi_logical_codec_from_pycapsule`. + /// datafusion-python invokes this with no arguments when the user + /// calls `ctx.with_logical_extension_codec(my_codec)`. The codec + /// owns its own bare `SessionContext` as a TaskContextProvider — + /// good enough for tests that only exercise UDF encode/decode. + fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let inner: Arc = Arc::new(CountingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec {}, + counters: Arc::clone(&self.counters), + }); + + let runtime = get_tokio_runtime().handle().clone(); + let bare_session: Arc = Arc::new(SessionContext::new()); + let ctx_provider = bare_session as Arc; + let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + + let name = cr"datafusion_logical_extension_codec".into(); + PyCapsule::new(py, ffi, Some(name)) + } +} diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs new file mode 100644 index 000000000..b1a586d9e --- /dev/null +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -0,0 +1,119 @@ +// 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. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::common::Result; +use datafusion::execution::{TaskContext, TaskContextProvider}; +use datafusion::logical_expr::ScalarUDF; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_proto::physical_plan::{DefaultPhysicalExtensionCodec, PhysicalExtensionCodec}; +use datafusion_python_util::get_tokio_runtime; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +#[derive(Debug, Default)] +pub(crate) struct PhysicalCallCounters { + pub encode_udf: AtomicUsize, + pub decode_udf: AtomicUsize, +} + +/// Mirror of [`super::logical_extension_codec::CountingLogicalExtensionCodec`] +/// for the physical layer. Delegates to `DefaultPhysicalExtensionCodec` +/// and bumps counters on UDF encode/decode so tests can prove the +/// session routed through a user-supplied physical codec. +#[derive(Debug)] +struct CountingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec, + counters: Arc, +} + +impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + ) -> Result> { + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + self.counters.decode_udf.fetch_add(1, Ordering::SeqCst); + self.inner.try_decode_udf(name, buf) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + self.counters.encode_udf.fetch_add(1, Ordering::SeqCst); + self.inner.try_encode_udf(node, buf) + } +} + +#[pyclass( + from_py_object, + name = "MyPhysicalExtensionCodec", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Clone)] +pub(crate) struct MyPhysicalExtensionCodec { + counters: Arc, +} + +#[pymethods] +impl MyPhysicalExtensionCodec { + #[new] + fn new() -> Self { + Self { + counters: Arc::new(PhysicalCallCounters::default()), + } + } + + fn encode_udf_calls(&self) -> usize { + self.counters.encode_udf.load(Ordering::SeqCst) + } + + fn decode_udf_calls(&self) -> usize { + self.counters.decode_udf.load(Ordering::SeqCst) + } + + fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let inner: Arc = + Arc::new(CountingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec {}, + counters: Arc::clone(&self.counters), + }); + + let runtime = get_tokio_runtime().handle().clone(); + let bare_session: Arc = Arc::new(SessionContext::new()); + let ctx_provider = bare_session as Arc; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + + let name = cr"datafusion_physical_extension_codec".into(); + PyCapsule::new(py, ffi, Some(name)) + } +} diff --git a/python/datafusion/context.py b/python/datafusion/context.py index dd6790402..5c3501941 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1750,4 +1750,22 @@ def with_logical_extension_codec(self, codec: Any) -> SessionContext: This only supports codecs that have been implemented using the FFI interface. """ - return self.ctx.with_logical_extension_codec(codec) + new_internal = self.ctx.with_logical_extension_codec(codec) + new = SessionContext.__new__(SessionContext) + new.ctx = new_internal + return new + + def __datafusion_physical_extension_codec__(self) -> Any: + """Access the PyCapsule FFI_PhysicalExtensionCodec.""" + return self.ctx.__datafusion_physical_extension_codec__() + + def with_physical_extension_codec(self, codec: Any) -> SessionContext: + """Create a new session context with the specified physical codec. + + This only supports codecs that have been implemented using the + FFI interface. + """ + new_internal = self.ctx.with_physical_extension_codec(codec) + new = SessionContext.__new__(SessionContext) + new.ctx = new_internal + return new diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 0f7f3ab5a..e0135e3ed 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -62,6 +62,7 @@ NullTreatment, RexType, ) + from datafusion.context import SessionContext from datafusion.plan import LogicalPlan @@ -432,6 +433,25 @@ def variant_name(self) -> str: """ return self.expr.variant_name() + def to_bytes(self, ctx: SessionContext | None = None) -> bytes: + """Serialize this expression to protobuf bytes. + + When ``ctx`` is supplied, encoding routes through the session's + installed :class:`LogicalExtensionCodec`. Without ``ctx`` a + default codec is used. + """ + ctx_arg = ctx.ctx if ctx is not None else None + return self.expr.to_bytes(ctx_arg) + + @staticmethod + def from_bytes(ctx: SessionContext, data: bytes) -> Expr: + """Decode an expression from serialized protobuf bytes. + + ``ctx`` provides the function registry for resolving UDF + references and the logical codec for in-band Python payloads. + """ + return Expr(expr_internal.RawExpr.from_bytes(ctx.ctx, data)) + def __richcmp__(self, other: Expr, op: int) -> Expr: """Comparison operator.""" return Expr(self.expr.__richcmp__(other.expr, op)) diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index c0cfd523f..b2c6eab3e 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -19,6 +19,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING, Any import datafusion._internal as df_internal @@ -88,19 +89,46 @@ def display_graphviz(self) -> str: return self._raw_plan.display_graphviz() @staticmethod - def from_proto(ctx: SessionContext, data: bytes) -> LogicalPlan: - """Create a LogicalPlan from protobuf bytes. + def from_bytes(ctx: SessionContext, data: bytes) -> LogicalPlan: + """Create a LogicalPlan from serialized protobuf bytes. - Tables created in memory from record batches are currently not supported. + Decoding routes through the session's installed + `LogicalExtensionCodec`. Tables created in memory from record + batches are currently not supported. """ - return LogicalPlan(df_internal.LogicalPlan.from_proto(ctx.ctx, data)) + return LogicalPlan(df_internal.LogicalPlan.from_bytes(ctx.ctx, data)) - def to_proto(self) -> bytes: - """Convert a LogicalPlan to protobuf bytes. + def to_bytes(self, ctx: SessionContext | None = None) -> bytes: + """Convert a LogicalPlan to serialized protobuf bytes. - Tables created in memory from record batches are currently not supported. + When ``ctx`` is supplied, encoding routes through the session's + installed `LogicalExtensionCodec` so user FFI codecs (registered + via :py:meth:`SessionContext.with_logical_extension_codec`) see + the encode path. With ``ctx=None`` a default codec is used. + Tables created in memory from record batches are currently not + supported. """ - return self._raw_plan.to_proto() + ctx_arg = ctx.ctx if ctx is not None else None + return self._raw_plan.to_bytes(ctx_arg) + + @staticmethod + def from_proto(ctx: SessionContext, data: bytes) -> LogicalPlan: + """Deprecated alias for :meth:`from_bytes`.""" + warnings.warn( + "LogicalPlan.from_proto is deprecated; use from_bytes instead", + DeprecationWarning, + stacklevel=2, + ) + return LogicalPlan.from_bytes(ctx, data) + + def to_proto(self) -> bytes: + """Deprecated alias for :meth:`to_bytes`.""" + warnings.warn( + "LogicalPlan.to_proto is deprecated; use to_bytes instead", + DeprecationWarning, + stacklevel=2, + ) + return self.to_bytes() def __eq__(self, other: LogicalPlan) -> bool: """Test equality.""" @@ -142,19 +170,43 @@ def partition_count(self) -> int: return self._raw_plan.partition_count @staticmethod - def from_proto(ctx: SessionContext, data: bytes) -> ExecutionPlan: - """Create an ExecutionPlan from protobuf bytes. + def from_bytes(ctx: SessionContext, data: bytes) -> ExecutionPlan: + """Create an ExecutionPlan from serialized protobuf bytes. - Tables created in memory from record batches are currently not supported. + Decoding routes through the session's installed + `PhysicalExtensionCodec`. Tables created in memory from record + batches are currently not supported. """ - return ExecutionPlan(df_internal.ExecutionPlan.from_proto(ctx.ctx, data)) + return ExecutionPlan(df_internal.ExecutionPlan.from_bytes(ctx.ctx, data)) - def to_proto(self) -> bytes: - """Convert an ExecutionPlan into protobuf bytes. + def to_bytes(self, ctx: SessionContext | None = None) -> bytes: + """Convert an ExecutionPlan into serialized protobuf bytes. - Tables created in memory from record batches are currently not supported. + When ``ctx`` is supplied, encoding routes through the session's + installed `PhysicalExtensionCodec`. Tables created in memory + from record batches are currently not supported. """ - return self._raw_plan.to_proto() + ctx_arg = ctx.ctx if ctx is not None else None + return self._raw_plan.to_bytes(ctx_arg) + + @staticmethod + def from_proto(ctx: SessionContext, data: bytes) -> ExecutionPlan: + """Deprecated alias for :meth:`from_bytes`.""" + warnings.warn( + "ExecutionPlan.from_proto is deprecated; use from_bytes instead", + DeprecationWarning, + stacklevel=2, + ) + return ExecutionPlan.from_bytes(ctx, data) + + def to_proto(self) -> bytes: + """Deprecated alias for :meth:`to_bytes`.""" + warnings.warn( + "ExecutionPlan.to_proto is deprecated; use to_bytes instead", + DeprecationWarning, + stacklevel=2, + ) + return self.to_bytes() def metrics(self) -> MetricsSet | None: """Return metrics for this plan node, or None if this plan has no MetricsSet. diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index 8aa791ae1..6a466f6f2 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -1178,3 +1178,29 @@ def test_round_trip_pyscalar_value(ctx: SessionContext, value: pa.Scalar): df = ctx.sql("select 1 as a") df = df.select(lit(value)) assert pa.table(df)[0][0] == value + + +def test_expr_to_bytes_roundtrip(ctx: SessionContext) -> None: + """An Expr round-trips through the session's logical codec.""" + from datafusion import Expr + + original = col("a") + lit(1) + blob = original.to_bytes(ctx) + restored = Expr.from_bytes(ctx, blob) + + # Canonical name preserves the structure of the expression even + # though the underlying PyExpr instances are different. + assert restored.canonical_name() == original.canonical_name() + + +def test_expr_to_bytes_no_ctx_default_codec() -> None: + """to_bytes(ctx=None) uses a default codec; builtin-only Exprs + still round-trip when a session is supplied on decode.""" + from datafusion import Expr + + fresh = SessionContext() + original = col("a") * lit(2) + blob = original.to_bytes() # encode side: default codec + restored = Expr.from_bytes(fresh, blob) + + assert restored.canonical_name() == original.canonical_name() diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 3705fc7ef..11e709f6b 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -35,21 +35,71 @@ def df(): return ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select("c1") -def test_logical_plan_to_proto(ctx, df) -> None: - logical_plan_bytes = df.logical_plan().to_proto() - logical_plan = LogicalPlan.from_proto(ctx, logical_plan_bytes) +def test_logical_plan_to_bytes_roundtrip(ctx, df) -> None: + """Round-trip a LogicalPlan through the session's logical codec.""" + logical_plan_bytes = df.logical_plan().to_bytes() + logical_plan = LogicalPlan.from_bytes(ctx, logical_plan_bytes) df_round_trip = ctx.create_dataframe_from_logical_plan(logical_plan) assert df.collect() == df_round_trip.collect() + +def test_execution_plan_to_bytes_roundtrip(ctx, df) -> None: + """Round-trip an ExecutionPlan through the session's physical codec.""" original_execution_plan = df.execution_plan() - execution_plan_bytes = original_execution_plan.to_proto() - execution_plan = ExecutionPlan.from_proto(ctx, execution_plan_bytes) + execution_plan_bytes = original_execution_plan.to_bytes() + execution_plan = ExecutionPlan.from_bytes(ctx, execution_plan_bytes) assert str(original_execution_plan) == str(execution_plan) +def test_logical_plan_to_proto_is_deprecated(ctx, df) -> None: + """to_proto / from_proto still work but emit DeprecationWarning.""" + plan = df.logical_plan() + + with pytest.warns(DeprecationWarning, match="to_proto"): + blob = plan.to_proto() + with pytest.warns(DeprecationWarning, match="from_proto"): + restored = LogicalPlan.from_proto(ctx, blob) + + df_round_trip = ctx.create_dataframe_from_logical_plan(restored) + assert df.collect() == df_round_trip.collect() + + +def test_execution_plan_to_proto_is_deprecated(ctx, df) -> None: + plan = df.execution_plan() + + with pytest.warns(DeprecationWarning, match="to_proto"): + blob = plan.to_proto() + with pytest.warns(DeprecationWarning, match="from_proto"): + restored = ExecutionPlan.from_proto(ctx, blob) + + assert str(plan) == str(restored) + + +def test_session_with_logical_extension_codec_roundtrip(ctx, df) -> None: + """A session with a non-default logical codec still round-trips builtins. + + The codec slot is overridable via with_logical_extension_codec; the + PythonLogicalCodec wrapper delegates unhandled cases to the inner + codec, so plans without Python UDFs are unaffected by the swap. + """ + # Default-routed session should round-trip via to_bytes. + blob = df.logical_plan().to_bytes() + restored = LogicalPlan.from_bytes(ctx, blob) + df_round_trip = ctx.create_dataframe_from_logical_plan(restored) + assert df.collect() == df_round_trip.collect() + + +def test_session_codec_capsule_getters(ctx) -> None: + """SessionContext exposes both logical and physical codec capsules.""" + logical = ctx.ctx.__datafusion_logical_extension_codec__() + physical = ctx.ctx.__datafusion_physical_extension_codec__() + assert logical is not None + assert physical is not None + + def test_metrics_tree_walk() -> None: ctx = SessionContext() ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')") From 8ba06e4147122962f356e11b73175379473f75a7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 18 May 2026 11:51:36 -0400 Subject: [PATCH 35/83] Update datafusion dependency to latest in preparation for DF54 (#1532) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: upgrade upstream DataFusion 53 → main (pre-54) Bump workspace deps to apache/datafusion@3d06bedc (git pin) in preparation for the 54.0.0 release. Workspace package version moves to 54.0.0 to track the upstream major convention. Compile fixes: - Drop as_any impls (trait now has Any as supertrait) and use the upstream-provided downcast_ref helper on dyn trait objects. - Reconcile FFI provider From conversions to drop redundant `+ Send` on Arc bounds. - Cast/TryCast: data_type → field.data_type() (FieldRef rename). - Stub match arms for new Expr::HigherOrderFunction / Lambda / LambdaVariable and ScalarValue::ListView / LargeListView variants; proper exposure deferred to PR 3 audit. - DatasetExec: partition_statistics returns Arc; add required apply_expressions trait method. - Suppress TableFunctionImpl::call deprecation pending call_with_args refactor that needs Session plumbing. User-facing test updates for upstream behavior changes: - median / approx_median / approx_percentile_cont now return Float64. - String functions (concat_ws, lower, upper, repeat, reverse, split_part, translate) return StringView when given StringView. - overlay appends past end-of-string rather than replacing the input. - arrays_zip / list_zip struct field names "c0"/"c1" → "1"/"2". - Filter on mismatched cast types now errors (was 0 matches). Co-Authored-By: Claude Opus 4.7 (1M context) * feat: expose DataFrame.alias and tidy public API after DF53→54 audit Companion to the upstream DataFusion 53 → main bump. The check-upstream audit (PR 3 of dev/release/upstream-sync.md) surfaced a small set of trivial wins; this commit ships them. Trivial wins: - DataFrame.alias(name) — wraps the logical plan in a SubqueryAlias. - functions.__all__: add `instr` and `position` (both were defined as public defs but missing from `__all__`, so they didn't show up in `from datafusion.functions import *` or generated docs). - top-level `datafusion.__all__`: re-export `TableProviderFactory` and `TableProviderFactoryExportable` (previously only reachable via the `datafusion.catalog` submodule). Non-trivial gaps surfaced by the audit (DataFrame.registry, into_*/task_ctx, SessionContext extensibility surface, distinct-aware aggregate variants, TableFunctionImpl::call_with_args migration, FFI Protocol pipeline gaps) are deferred — each warrants its own design and PR. Co-Authored-By: Claude Opus 4.7 (1M context) * taplo fmt * Update unit test to go along with https://github.com/apache/datafusion/pull/22133 * docs: demonstrate alias via self-join in DataFrame.alias example Prior example called alias("t") then to_pydict(), which did not show the qualifier effect. Replace with a self-join that uses col("l.val") and col("r.val") so the disambiguation behavior is visible. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: wrap higher-order, lambda, and lambda-variable Expr variants DataFusion 54 introduces Expr::HigherOrderFunction, Expr::Lambda, and Expr::LambdaVariable. PyExpr::to_variant previously errored on each with py_unsupported_variant_err. Add PyHigherOrderFunction, PyLambda, and PyLambdaVariable wrappers, register them in the expr pymodule and re-export from python/datafusion/expr.py, and dispatch to_variant to the new wrappers. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: wire rex_type and rex_call_operands for new Expr variants Map HigherOrderFunction and Lambda to RexType::Call; LambdaVariable to RexType::Reference. In rex_call_operands return the args for HigherOrderFunction, the body for Lambda, and self for LambdaVariable (mirroring Column). In rex_call_operator return the underlying UDF name for HigherOrderFunction and the literal "lambda" for Lambda. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: support LargeList/ListView/LargeListView in map_from_scalar_to_arrow These ScalarValue variants all wrap Arc<...Array>, exposing the outer DataType via Array::data_type(), so we can mirror the existing ScalarValue::List arm instead of returning PyNotImplementedError. This makes Expr.types() work for plans that round-trip through SQL or proto where these scalar variants surface. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: switch PyTableFunction to non-deprecated call_with_args DataFusion 53.0.0 deprecated TableFunctionImpl::call in favor of call_with_args(args: TableFunctionArgs), which threads a Session reference alongside the exprs. Implement call_with_args on PyTableFunction (delegating to the FFI variant's call_with_args, or ignoring the session for the pure-Python variant which doesn't use it) and have __call__ build a TableFunctionArgs from the global session. Drops both #[allow(deprecated)] attributes. Co-Authored-By: Claude Opus 4.7 (1M context) * build: revert workspace version to 53.0.0 and move DF overrides to [patch.crates-io] The workspace version was prematurely bumped to 54.0.0 in the DF53→pre-54 upgrade. Restore it to 53.0.0 until we are actually ready to cut the 54 release. The same change had moved every datafusion-* dependency from a crates.io version constraint to a direct git dep in [workspace.dependencies]. Switch them back to "version = \"53\"" and move the git rev overrides into [patch.crates-io] so the published manifest will be patch-free. Co-Authored-By: Claude Opus 4.7 (1M context) * taplo format * test: sort FFI test results by partition key before equality compare Multi-partition `collect()` returns batches in execution-scheduling order, which is non-deterministic and differs between local and CI runners. Sort by the first value of column 0 (unique per partition in each affected test) so the expected/actual comparison is stable. Co-Authored-By: Claude Opus 4.7 (1M context) * Bump datafusion main commit * test: cover new DF54 expr wrappers, catalog factories, and DataFrame.alias Add module-metadata checks for HigherOrderFunction, Lambda, LambdaVariable and the top-level TableProviderFactory / TableProviderFactoryExportable re-exports, plus a self-join regression test exercising the new DataFrame.alias() qualifier-based selection path. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .ai/skills/check-upstream/SKILL.md | 45 + Cargo.lock | 827 ++++++++---------- Cargo.toml | 9 + crates/core/src/catalog.rs | 53 +- crates/core/src/common/data_type.rs | 6 +- crates/core/src/common/schema.rs | 5 - crates/core/src/context.rs | 13 +- crates/core/src/dataframe.rs | 5 + crates/core/src/dataset.rs | 7 - crates/core/src/dataset_exec.rs | 20 +- crates/core/src/expr.rs | 42 +- crates/core/src/expr/cast.rs | 4 +- crates/core/src/expr/higher_order_function.rs | 84 ++ crates/core/src/expr/lambda.rs | 78 ++ crates/core/src/expr/lambda_variable.rs | 67 ++ crates/core/src/table.rs | 5 - crates/core/src/udf.rs | 5 - crates/core/src/udtf.rs | 16 +- crates/core/src/udwf.rs | 5 - .../python/tests/_test_schema_provider.py | 7 +- .../python/tests/_test_table_function.py | 7 +- .../python/tests/_test_table_provider.py | 6 +- .../src/aggregate_udf.rs | 5 - .../src/catalog_provider.rs | 13 - .../datafusion-ffi-example/src/scalar_udf.rs | 5 - .../datafusion-ffi-example/src/window_udf.rs | 5 - python/datafusion/__init__.py | 9 +- python/datafusion/dataframe.py | 26 + python/datafusion/expr.py | 6 + python/datafusion/functions.py | 4 +- python/tests/test_aggregation.py | 34 +- python/tests/test_dataframe.py | 14 + python/tests/test_expr.py | 11 +- python/tests/test_functions.py | 27 +- python/tests/test_imports.py | 12 + 35 files changed, 865 insertions(+), 622 deletions(-) create mode 100644 crates/core/src/expr/higher_order_function.rs create mode 100644 crates/core/src/expr/lambda.rs create mode 100644 crates/core/src/expr/lambda_variable.rs diff --git a/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md index ac4835a4e..3bac018ef 100644 --- a/.ai/skills/check-upstream/SKILL.md +++ b/.ai/skills/check-upstream/SKILL.md @@ -29,6 +29,29 @@ You are auditing the datafusion-python project to find features from the upstrea **IMPORTANT: The Python API is the source of truth for coverage.** A function or method is considered "exposed" if it exists in the Python API (e.g., `python/datafusion/functions.py`), even if there is no corresponding entry in the Rust bindings. Many upstream functions are aliases of other functions — the Python layer can expose these aliases by calling a different underlying Rust binding. Do NOT report a function as missing if it appears in the Python `__all__` list and has a working implementation, regardless of whether a matching `#[pyfunction]` exists in Rust. +**IMPORTANT: audit the total upstream surface, not the delta since the last pin.** Gaps accumulate across syncs. A patch-release bump with a "bug fixes only" changelog does not mean there is nothing to find — pre-existing gaps from earlier majors still need to be surfaced. Always run the full comparison. + +## Compile-Signal Triggers + +If a recent upstream bump required *any* of the following while fixing +compile errors in `crates/core/` or the FFI example, treat that as a +**hard signal** that user-facing surface area grew and run this skill +before considering the bump done. Each pattern corresponds to a class of +gap that frequently shows up in the audit: + +| Signal during PR 1 compile fix | Likely gap to check | +|---|---| +| New `Expr::*` variant added to a non-exhaustive `match` (`HigherOrderFunction`, `Lambda`, `LambdaVariable`, …) | New lambda / higher-order scalar functions (`any_match`, `array_transform`, `list_transform`, …) | +| New `ScalarValue::*` variant (`ListView`, `LargeListView`, …) | New scalar / array functions that consume or produce the type | +| New required trait method on `ExecutionPlan` / `TableProvider` / `*UDFImpl` (`apply_expressions`, …) | Corresponding capability on the Python wrapper class | +| Renamed or restructured struct field (e.g. `Cast.data_type` → `Cast.field: FieldRef`) | Any Python accessor / SKILL.md doc that read the old field | +| Newly deprecated trait method with a `_with_args` / `_with_options` replacement | The `*_with_options` variant frequently warrants a separate Python entry point | + +PR 1 of `dev/release/upstream-sync.md` asks you to log these signals as +they appear. When you run this skill, use that log as a checklist: every +entry must either show up in the audit output or be explicitly skipped +with a reason. + ## Areas to Check The user may specify an area via `$ARGUMENTS`. If no area is specified or "all" is given, check all areas. @@ -173,6 +196,28 @@ These upstream FFI types have been reviewed and do not need to be independently - FFI example in `examples/datafusion-ffi-example/` - Type appears in union type hints where accepted +### 8. `__all__` Hygiene (functions.py) + +Independent of upstream parity, also flag public `def` symbols in +`python/datafusion/functions.py` that are missing from the module's +`__all__`. These are functions a user can call but that do not show up in +`from datafusion.functions import *`, in tab-completion against the +namespace, or in generated API docs — typically an oversight rather than +an intentional omission. + +**How to check:** +1. Grep for `^def ([a-z_][a-z0-9_]*)\(` in `python/datafusion/functions.py` + to enumerate every public function definition. +2. Read the `__all__` list at the top of the same file. +3. Report any function in (1) that is not in (2). Skip private helpers + (names starting with `_`). + +A historical example: `instr` and `position` shipped as public `def`s but +were absent from `__all__` until the gap was caught here. + +For each finding, propose adding the name to `__all__` in alphabetical +position with the existing entries. + ## Checking for Existing GitHub Issues After identifying missing APIs, search the open issues at https://github.com/apache/datafusion-python/issues for each gap to see if an issue already exists requesting that API be exposed. Search using the function or method name as the query. diff --git a/Cargo.lock b/Cargo.lock index 1d148b0e1..0c4b77582 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,54 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "abi_stable" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69d6512d3eb05ffe5004c59c206de7f99c34951504056ce23fc953842f12c445" -dependencies = [ - "abi_stable_derive", - "abi_stable_shared", - "const_panic", - "core_extensions", - "crossbeam-channel", - "generational-arena", - "libloading", - "lock_api", - "parking_lot", - "paste", - "repr_offset", - "rustc_version", - "serde", - "serde_derive", - "serde_json", -] - -[[package]] -name = "abi_stable_derive" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7178468b407a4ee10e881bc7a328a65e739f0863615cca4429d43916b05e898" -dependencies = [ - "abi_stable_shared", - "as_derive_utils", - "core_extensions", - "proc-macro2", - "quote", - "rustc_version", - "syn 1.0.109", - "typed-arena", -] - -[[package]] -name = "abi_stable_shared" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2b5df7688c123e63f4d4d649cba63f2967ba7f7861b1664fca3f77d3dad2b63" -dependencies = [ - "core_extensions", -] - [[package]] name = "adler2" version = "2.0.1" @@ -115,35 +67,6 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" -[[package]] -name = "apache-avro" -version = "0.21.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36fa98bc79671c7981272d91a8753a928ff6a1cd8e4f20a44c45bd5d313840bf" -dependencies = [ - "bigdecimal", - "bon", - "bzip2", - "crc32fast", - "digest", - "liblzma", - "log", - "miniz_oxide", - "num-bigint", - "quad-rand", - "rand 0.9.4", - "regex-lite", - "serde", - "serde_bytes", - "serde_json", - "snap", - "strum", - "strum_macros", - "thiserror", - "uuid", - "zstd", -] - [[package]] name = "ar_archive_writer" version = "0.5.1" @@ -176,9 +99,9 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d441fdda254b65f3e9025910eb2c2066b6295d9c8ed409522b8d2ace1ff8574c" +checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" dependencies = [ "arrow-arith", "arrow-array", @@ -198,9 +121,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ced5406f8b720cc0bc3aa9cf5758f93e8593cda5490677aa194e4b4b383f9a59" +checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" dependencies = [ "arrow-array", "arrow-buffer", @@ -229,6 +152,30 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrow-avro" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "049230728cd6e093088c8d231b4beede184e35cad7777c1505c0d5a8571f4376" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-schema", + "bytes", + "bzip2", + "crc", + "flate2", + "indexmap", + "liblzma", + "rand 0.9.2", + "serde", + "serde_json", + "snap", + "strum_macros", + "uuid", + "zstd", +] + [[package]] name = "arrow-buffer" version = "58.3.0" @@ -243,9 +190,9 @@ dependencies = [ [[package]] name = "arrow-cast" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0127816c96533d20fc938729f48c52d3e48f99717e7a0b5ade77d742510736d" +checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" dependencies = [ "arrow-array", "arrow-buffer", @@ -265,9 +212,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca025bd0f38eeecb57c2153c0123b960494138e6a957bbda10da2b25415209fe" +checksum = "e94e8cf7e517657a52b91ea1263acf38c4ca62a84655d72458a3359b12ab97de" dependencies = [ "arrow-array", "arrow-cast", @@ -293,9 +240,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "609a441080e338147a84e8e6904b6da482cefb957c5cdc0f3398872f69a315d0" +checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" dependencies = [ "arrow-array", "arrow-buffer", @@ -309,15 +256,16 @@ dependencies = [ [[package]] name = "arrow-json" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ead0914e4861a531be48fe05858265cf854a4880b9ed12618b1d08cba9bebc8" +checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" dependencies = [ "arrow-array", "arrow-buffer", "arrow-cast", - "arrow-data", + "arrow-ord", "arrow-schema", + "arrow-select", "chrono", "half", "indexmap", @@ -333,9 +281,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "763a7ba279b20b52dad300e68cfc37c17efa65e68623169076855b3a9e941ca5" +checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" dependencies = [ "arrow-array", "arrow-buffer", @@ -346,9 +294,9 @@ dependencies = [ [[package]] name = "arrow-pyarrow" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e63351dc11981a316c828a6032a5021345bba882f68bc4a36c36825a50725089" +checksum = "d29abdf672a81c1aeb57fd2661457f9918964d49aed0e9f18932535f2a9e49ce" dependencies = [ "arrow-array", "arrow-data", @@ -358,9 +306,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e14fe367802f16d7668163ff647830258e6e0aeea9a4d79aaedf273af3bdcd3e" +checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" dependencies = [ "arrow-array", "arrow-buffer", @@ -382,9 +330,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78694888660a9e8ac949853db393af2a8b8fc82c19ce333132dfa2e72cc1a7fe" +checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" dependencies = [ "ahash", "arrow-array", @@ -396,9 +344,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61e04a01f8bb73ce54437514c5fd3ee2aa3e8abe4c777ee5cc55853b1652f79e" +checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" dependencies = [ "arrow-array", "arrow-buffer", @@ -411,18 +359,6 @@ dependencies = [ "regex-syntax", ] -[[package]] -name = "as_derive_utils" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff3c96645900a44cf11941c111bd08a6573b0e2f9f69bc9264b179d8fae753c4" -dependencies = [ - "core_extensions", - "proc-macro2", - "quote", - "syn 1.0.109", -] - [[package]] name = "async-compression" version = "0.4.41" @@ -440,9 +376,6 @@ name = "async-ffi" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f4de21c0feef7e5a556e51af767c953f0501f7f300ba785cc99c47bdc8081a50" -dependencies = [ - "abi_stable", -] [[package]] name = "async-recursion" @@ -504,7 +437,6 @@ dependencies = [ "num-bigint", "num-integer", "num-traits", - "serde", ] [[package]] @@ -519,7 +451,7 @@ version = "0.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" dependencies = [ - "digest", + "digest 0.10.7", ] [[package]] @@ -546,28 +478,12 @@ dependencies = [ ] [[package]] -name = "bon" -version = "3.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f47dbe92550676ee653353c310dfb9cf6ba17ee70396e1f7cf0a2020ad49b2fe" -dependencies = [ - "bon-macros", - "rustversion", -] - -[[package]] -name = "bon-macros" -version = "3.9.1" +name = "block-buffer" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "519bd3116aeeb42d5372c29d982d16d0170d3d4a5ed85fc7dd91642ffff3c67c" +checksum = "cdd35008169921d80bc60d3d0ab416eecb028c4cd653352907921d95084790be" dependencies = [ - "darling", - "ident_case", - "prettyplease", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.117", + "hybrid-array", ] [[package]] @@ -715,6 +631,12 @@ version = "0.4.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75984efb6ed102a0d42db99afb6c1948f0380d1d91808d5529916e6c08b49d8d" +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "const-random" version = "0.1.18" @@ -735,15 +657,6 @@ dependencies = [ "tiny-keccak", ] -[[package]] -name = "const_panic" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e262cdaac42494e3ae34c43969f9cdeb7da178bdb4b66fa6a1ea2edb4c8ae652" -dependencies = [ - "typewit", -] - [[package]] name = "constant_time_eq" version = "0.4.2" @@ -766,21 +679,6 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "core_extensions" -version = "1.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42bb5e5d0269fd4f739ea6cedaf29c16d81c27a7ce7582008e90eb50dcd57003" -dependencies = [ - "core_extensions_proc_macros", -] - -[[package]] -name = "core_extensions_proc_macros" -version = "1.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "533d38ecd2709b7608fb8e18e4504deb99e9a72879e6aa66373a76d8dc4259ea" - [[package]] name = "cpufeatures" version = "0.2.17" @@ -800,21 +698,27 @@ dependencies = [ ] [[package]] -name = "crc32fast" -version = "1.5.0" +name = "crc" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" dependencies = [ - "cfg-if", + "crc-catalog", ] [[package]] -name = "crossbeam-channel" -version = "0.5.15" +name = "crc-catalog" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" dependencies = [ - "crossbeam-utils", + "cfg-if", ] [[package]] @@ -839,6 +743,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710" +dependencies = [ + "hybrid-array", +] + [[package]] name = "cstr" version = "0.2.12" @@ -870,40 +783,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "darling" -version = "0.23.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" -dependencies = [ - "darling_core", - "darling_macro", -] - -[[package]] -name = "darling_core" -version = "0.23.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" -dependencies = [ - "ident_case", - "proc-macro2", - "quote", - "strsim", - "syn 2.0.117", -] - -[[package]] -name = "darling_macro" -version = "0.23.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" -dependencies = [ - "darling_core", - "quote", - "syn 2.0.117", -] - [[package]] name = "dashmap" version = "6.1.0" @@ -921,13 +800,11 @@ dependencies = [ [[package]] name = "datafusion" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93db0e623840612f7f2cd757f7e8a8922064192363732c88692e0870016e141b" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "arrow-schema", "async-trait", - "bytes", "bzip2", "chrono", "datafusion-catalog", @@ -958,14 +835,13 @@ dependencies = [ "datafusion-sql", "flate2", "futures", + "indexmap", "itertools", "liblzma", "log", "object_store", "parking_lot", "parquet", - "rand 0.9.4", - "regex", "sqlparser", "tempfile", "tokio", @@ -977,8 +853,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37cefde60b26a7f4ff61e9d2ff2833322f91df2b568d7238afe67bde5bdffb66" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "async-trait", @@ -1002,8 +877,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17e112307715d6a7a331111a4c2330ff54bc237183511c319e3708a4cff431fb" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "async-trait", @@ -1025,34 +899,32 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d72a11ca44a95e1081870d3abb80c717496e8a7acb467a1d3e932bb636af5cc2" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "ahash", - "apache-avro", "arrow", "arrow-ipc", + "arrow-schema", "chrono", + "foldhash 0.2.0", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "indexmap", "itertools", "libc", "log", "object_store", "parquet", - "paste", "recursive", "sqlparser", "tokio", + "uuid", "web-time", ] [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89f4afaed29670ec4fd6053643adc749fe3f4bc9d1ce1b8c5679b22c67d12def" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "futures", "log", @@ -1062,8 +934,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9fb386e1691355355a96419978a0022b7947b44d4a24a6ea99f00b6b485cbb6" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "async-compression", @@ -1087,7 +958,8 @@ dependencies = [ "liblzma", "log", "object_store", - "rand 0.9.4", + "parking_lot", + "rand 0.9.2", "tokio", "tokio-util", "url", @@ -1097,8 +969,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffa6c52cfed0734c5f93754d1c0175f558175248bf686c944fb05c373e5fc096" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "arrow-ipc", @@ -1121,28 +992,25 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a579c3bd290c66ea4b269493e75e8a3ed42c9c895a651f10210a29538aee50c4" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "apache-avro", "arrow", + "arrow-avro", "async-trait", "bytes", "datafusion-common", "datafusion-datasource", - "datafusion-physical-expr-common", + "datafusion-physical-expr-adapter", "datafusion-physical-plan", "datafusion-session", "futures", - "num-traits", "object_store", ] [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "503f29e0582c1fc189578d665ff57d9300da1f80c282777d7eb67bb79fb8cdca" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "async-trait", @@ -1164,8 +1032,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e33804749abc8d0c8cb7473228483cb8070e524c6f6086ee1b85a64debe2b3d2" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "async-trait", @@ -1180,7 +1047,6 @@ dependencies = [ "datafusion-session", "futures", "object_store", - "serde_json", "tokio", "tokio-stream", ] @@ -1188,8 +1054,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a8e0365e0e08e8ff94d912f0ababcf9065a1a304018ba90b1fc83c855b4997" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "async-trait", @@ -1199,6 +1064,7 @@ dependencies = [ "datafusion-datasource", "datafusion-execution", "datafusion-expr", + "datafusion-functions", "datafusion-functions-aggregate-common", "datafusion-physical-expr", "datafusion-physical-expr-adapter", @@ -1218,19 +1084,16 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8de6ac0df1662b9148ad3c987978b32cbec7c772f199b1d53520c8fa764a87ee" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c03c7fbdaefcca4ef6ffe425a5fc2325763bfb426599bb0bf4536466efabe709" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "arrow-buffer", "async-trait", - "chrono", "dashmap", "datafusion-common", "datafusion-expr", @@ -1239,7 +1102,7 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.4", + "rand 0.9.2", "tempfile", "url", ] @@ -1247,10 +1110,10 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "574b9b6977fedbd2a611cbff12e5caf90f31640ad9dc5870f152836d94bad0dd" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", + "arrow-schema", "async-trait", "chrono", "datafusion-common", @@ -1261,7 +1124,6 @@ dependencies = [ "datafusion-physical-expr-common", "indexmap", "itertools", - "paste", "recursive", "serde_json", "sqlparser", @@ -1270,27 +1132,24 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d7c3adf3db8bf61e92eb90cb659c8e8b734593a8f7c8e12a843c7ddba24b87e" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "datafusion-common", "indexmap", "itertools", - "paste", ] [[package]] name = "datafusion-ffi" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b95173344d04ba62755c949bf44f8d1a6e4414cf6392a635db96c07e711b9a3c" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "abi_stable", "arrow", "arrow-schema", "async-ffi", "async-trait", + "chrono", "datafusion-catalog", "datafusion-common", "datafusion-datasource", @@ -1299,14 +1158,17 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-physical-expr", "datafusion-physical-expr-common", + "datafusion-physical-optimizer", "datafusion-physical-plan", "datafusion-proto", "datafusion-proto-common", "datafusion-session", "futures", + "libloading", "log", "prost", "semver", + "stabby", "tokio", ] @@ -1335,8 +1197,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f28aa4e10384e782774b10e72aca4d93ef7b31aa653095d9d4536b0a3dbc51b6" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "arrow-buffer", @@ -1351,26 +1212,24 @@ dependencies = [ "datafusion-expr", "datafusion-expr-common", "datafusion-macros", + "datafusion-physical-expr-common", "hex", "itertools", "log", - "md-5", + "md-5 0.11.0", "memchr", "num-traits", - "rand 0.9.4", + "rand 0.9.2", "regex", "sha2", - "unicode-segmentation", "uuid", ] [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00aa6217e56098ba84e0a338176fe52f0a84cca398021512c6c8c5eff806d0ad" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-doc", @@ -1380,19 +1239,17 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr", "datafusion-physical-expr-common", + "foldhash 0.2.0", "half", "log", "num-traits", - "paste", ] [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b511250349407db7c43832ab2de63f5557b19a20dfd236b39ca2c04468b50d47" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr-common", @@ -1402,8 +1259,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef13a858e20d50f0a9bb5e96e7ac82b4e7597f247515bccca4fdd2992df0212a" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "arrow-ord", @@ -1417,18 +1273,17 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-macros", "datafusion-physical-expr-common", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "itertools", "itoa", "log", - "paste", + "memchr", ] [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b40d3f5bbb3905f9ccb1ce9485a9595c77b69758a7c24d3ba79e334ff51e7e" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "async-trait", @@ -1437,14 +1292,12 @@ dependencies = [ "datafusion-expr", "datafusion-physical-plan", "parking_lot", - "paste", ] [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4e88ec9d57c9b685d02f58bfee7be62d72610430ddcedb82a08e5d9925dbfb6" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "datafusion-common", @@ -1455,14 +1308,12 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "log", - "paste", ] [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8307bb93519b1a91913723a1130cfafeee3f72200d870d88e91a6fc5470ede5c" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1471,8 +1322,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e367e6a71051d0ebdd29b2f85d12059b38b1d1f172c6906e80016da662226bd" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "datafusion-doc", "quote", @@ -1482,8 +1332,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e929015451a67f77d9d8b727b2bf3a40c4445fdef6cdc53281d7d97c76888ace" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "chrono", @@ -1502,10 +1351,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b1e68aba7a4b350401cfdf25a3d6f989ad898a7410164afe9ca52080244cb59" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr", @@ -1513,11 +1360,10 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "indexmap", "itertools", "parking_lot", - "paste", "petgraph", "recursive", "tokio", @@ -1526,8 +1372,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea22315f33cf2e0adc104e8ec42e285f6ed93998d565c65e82fec6a9ee9f9db4" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "datafusion-common", @@ -1541,25 +1386,23 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b04b45ea8ad3ac2d78f2ea2a76053e06591c9629c7a603eda16c10649ecf4362" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "ahash", "arrow", "chrono", "datafusion-common", "datafusion-expr-common", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "indexmap", "itertools", "parking_lot", + "pin-project", ] [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7cb13397809a425918f608dfe8653f332015a3e330004ab191b4404187238b95" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "datafusion-common", @@ -1577,11 +1420,11 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5edc023675791af9d5fb4cc4c24abf5f7bd3bd4dcf9e5bd90ea1eff6976dcc79" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ - "ahash", "arrow", + "arrow-data", + "arrow-ipc", "arrow-ord", "arrow-schema", "async-trait", @@ -1596,7 +1439,7 @@ dependencies = [ "datafusion-physical-expr-common", "futures", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "indexmap", "itertools", "log", @@ -1609,8 +1452,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a387aaef949dc16bb6abc81bd1af850ec7449183aef011214f9724957495738" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "chrono", @@ -1631,14 +1473,12 @@ dependencies = [ "datafusion-proto-common", "object_store", "prost", - "rand 0.9.4", ] [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16e614c7c53a9c304c6a850b821010bb492e57300311835f1180613f9d2c63d9" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "datafusion-common", @@ -1648,8 +1488,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac8c76860e355616555081cab5968cec1af7a80701ff374510860bcd567e365a" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "datafusion-common", @@ -1658,7 +1497,6 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", - "itertools", "log", ] @@ -1709,8 +1547,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5412111aa48e2424ba926112e192f7a6b7e4ccb450145d25ce5ede9f19dc491e" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "async-trait", "datafusion-common", @@ -1723,8 +1560,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa0d133ddf8b9b3b872acac900157f783e7b879fe9a6bccf389abebbfac45ec1" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "arrow", "bigdecimal", @@ -1741,9 +1577,8 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "53.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5e5656a7e63d51dd3e5af3dbd347ea83bbe993a77c66b854b74961570d16490" +version = "53.1.0" +source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" dependencies = [ "async-recursion", "async-trait", @@ -1765,11 +1600,22 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", "subtle", ] +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.0", + "const-oid", + "crypto-common 0.2.1", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -1963,15 +1809,6 @@ dependencies = [ "slab", ] -[[package]] -name = "generational-arena" -version = "0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "877e94aff08e743b651baaea359664321055749b398adff8740a7399af7796e7" -dependencies = [ - "cfg-if", -] - [[package]] name = "generic-array" version = "0.14.7" @@ -2091,6 +1928,11 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] [[package]] name = "heck" @@ -2149,6 +1991,15 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "hybrid-array" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.8.1" @@ -2322,12 +2173,6 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" -[[package]] -name = "ident_case" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" - [[package]] name = "idna" version = "1.1.0" @@ -2351,12 +2196,12 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.13.0" +version = "2.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "serde", "serde_core", ] @@ -2489,18 +2334,18 @@ checksum = "2c4a545a15244c7d945065b5d392b2d2d7f21526fba56ce51467b06ed445e8f7" [[package]] name = "libc" -version = "0.2.183" +version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libloading" -version = "0.7.4" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" dependencies = [ "cfg-if", - "winapi", + "windows-link", ] [[package]] @@ -2531,11 +2376,12 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.47" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d1eacfa31c33ec25e873c136ba5669f00f9866d0688bea7be4d3f7e43067df6" +checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" dependencies = [ "cc", + "libc", ] [[package]] @@ -2587,7 +2433,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" dependencies = [ "cfg-if", - "digest", + "digest 0.10.7", +] + +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", ] [[package]] @@ -2598,9 +2454,9 @@ checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "mimalloc" -version = "0.1.50" +version = "0.1.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3627c4272df786b9260cabaa46aec1d59c93ede723d4c3ef646c503816b0640" +checksum = "e1ee66a4b64c74f4ef288bcbb9192ad9c3feaad75193129ac8509af543894fd8" dependencies = [ "libmimalloc-sys", ] @@ -2640,7 +2496,6 @@ checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" dependencies = [ "num-integer", "num-traits", - "serde", ] [[package]] @@ -2700,7 +2555,7 @@ dependencies = [ "humantime", "hyper", "itertools", - "md-5", + "md-5 0.10.6", "parking_lot", "percent-encoding", "quick-xml", @@ -2766,9 +2621,9 @@ dependencies = [ [[package]] name = "parquet" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d3f9f2205199603564127932b89695f52b62322f541d0fc7179d57c2e1c9877" +checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" dependencies = [ "ahash", "arrow-array", @@ -2784,7 +2639,7 @@ dependencies = [ "flate2", "futures", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "lz4_flex", "num-bigint", "num-integer", @@ -2879,6 +2734,26 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project" +version = "1.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbf0d9e68100b3a7989b4901972f265cd542e560a3a8a724e1e20322f4d06ce9" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a990e22f43e84855daf260dded30524ef4a9021cc7541c26540500a50b624389" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -2931,6 +2806,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -3093,12 +2977,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "quad-rand" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a651516ddc9168ebd67b24afd085a718be02f8858fe406591b013d101ce2f40" - [[package]] name = "quick-xml" version = "0.39.2" @@ -3138,7 +3016,7 @@ dependencies = [ "bytes", "getrandom 0.3.4", "lru-slab", - "rand 0.9.4", + "rand 0.9.2", "ring", "rustc-hash", "rustls", @@ -3187,11 +3065,22 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.9.4" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" dependencies = [ - "rand_chacha", + "rand_chacha 0.9.0", "rand_core 0.9.5", ] @@ -3206,6 +3095,16 @@ dependencies = [ "rand_core 0.10.0", ] +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + [[package]] name = "rand_chacha" version = "0.9.0" @@ -3216,6 +3115,15 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + [[package]] name = "rand_core" version = "0.9.5" @@ -3283,12 +3191,6 @@ dependencies = [ "regex-syntax", ] -[[package]] -name = "regex-lite" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" - [[package]] name = "regex-syntax" version = "0.8.10" @@ -3305,15 +3207,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "repr_offset" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb1070755bd29dffc19d0971cab794e607839ba2ef4b69a9e6fbc8733c1b72ea" -dependencies = [ - "tstr", -] - [[package]] name = "reqwest" version = "0.12.28" @@ -3436,9 +3329,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "df33b2b81ac578cabaf06b89b0631153a3f416b0a886e8a7a1707fb51abbd1ef" dependencies = [ "ring", "rustls-pki-types", @@ -3530,9 +3423,9 @@ dependencies = [ [[package]] name = "semver" -version = "1.0.27" +version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" dependencies = [ "serde", "serde_core", @@ -3554,16 +3447,6 @@ dependencies = [ "serde_derive", ] -[[package]] -name = "serde_bytes" -version = "0.11.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" -dependencies = [ - "serde", - "serde_core", -] - [[package]] name = "serde_core" version = "1.0.228" @@ -3601,6 +3484,7 @@ version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ + "indexmap", "itoa", "memchr", "serde", @@ -3647,15 +3531,21 @@ dependencies = [ [[package]] name = "sha2" -version = "0.10.9" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", - "cpufeatures 0.2.17", - "digest", + "cpufeatures 0.3.0", + "digest 0.11.3", ] +[[package]] +name = "sha2-const-stable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f179d4e11094a893b82fff208f74d448a7512f99f5a0acbd5c679b705f83ed9" + [[package]] name = "shlex" version = "1.3.0" @@ -3710,9 +3600,9 @@ dependencies = [ [[package]] name = "sqlparser" -version = "0.61.0" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbf5ea8d4d7c808e1af1cbabebca9a2abe603bcefc22294c5b95018d53200cb7" +checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec" dependencies = [ "log", "recursive", @@ -3730,6 +3620,41 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "stabby" +version = "72.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "976399a0c48ea769ef7f5dc303bb88240ab8d84008647a6b2303eced3dab3945" +dependencies = [ + "rustversion", + "stabby-abi", +] + +[[package]] +name = "stabby-abi" +version = "72.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7b54832a9a1f92a0e55e74a5c0332744426edc515bb3fbad82f10b874a87f0d" +dependencies = [ + "rustc_version", + "rustversion", + "sha2-const-stable", + "stabby-macros", +] + +[[package]] +name = "stabby-macros" +version = "72.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a768b1e51e4dbfa4fa52ae5c01241c0a41e2938fdffbb84add0c8238092f9091" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "rand 0.8.6", + "syn 1.0.109", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -3749,23 +3674,11 @@ dependencies = [ "windows-sys 0.59.0", ] -[[package]] -name = "strsim" -version = "0.11.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" - -[[package]] -name = "strum" -version = "0.27.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" - [[package]] name = "strum_macros" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" dependencies = [ "heck", "proc-macro2", @@ -3775,11 +3688,12 @@ dependencies = [ [[package]] name = "substrait" -version = "0.62.2" +version = "0.63.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62fc4b483a129b9772ccb9c3f7945a472112fdd9140da87f8a4e7f1d44e045d0" +checksum = "e620ff4d5c02fd6f7752931aa74b16a26af66a63022cc1ad412c77edbe0bab47" dependencies = [ "heck", + "indexmap", "pbjson", "pbjson-build", "pbjson-types", @@ -3992,6 +3906,36 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.11+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow", +] + [[package]] name = "tower" version = "0.5.3" @@ -4074,44 +4018,17 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" -[[package]] -name = "tstr" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f8e0294f14baae476d0dd0a2d780b2e24d66e349a9de876f5126777a37bdba7" -dependencies = [ - "tstr_proc_macros", -] - -[[package]] -name = "tstr_proc_macros" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e78122066b0cb818b8afd08f7ed22f7fdbc3e90815035726f0840d0d26c0747a" - [[package]] name = "twox-hash" version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" -[[package]] -name = "typed-arena" -version = "2.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" - [[package]] name = "typenum" -version = "1.19.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" - -[[package]] -name = "typewit" -version = "1.14.2" +version = "1.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8c1ae7cc0fdb8b842d65d127cb981574b0d2b249b74d1c7a2986863dc134f71" +checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" [[package]] name = "typify" @@ -4216,13 +4133,12 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.23.1" +version = "1.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" +checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9" dependencies = [ "getrandom 0.4.2", "js-sys", - "serde_core", "wasm-bindgen", ] @@ -4401,22 +4317,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "winapi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - [[package]] name = "winapi-util" version = "0.1.11" @@ -4426,12 +4326,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - [[package]] name = "windows-core" version = "0.62.2" @@ -4656,6 +4550,15 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +[[package]] +name = "winnow" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee1708bef14716a11bae175f579062d4554d95be2c6829f518df847b7b3fdd0" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.51.0" diff --git a/Cargo.toml b/Cargo.toml index 077bc093f..13d7040a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,3 +71,12 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] +datafusion = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-common = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } diff --git a/crates/core/src/catalog.rs b/crates/core/src/catalog.rs index 30ec4744c..8ad49b098 100644 --- a/crates/core/src/catalog.rs +++ b/crates/core/src/catalog.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::collections::HashSet; use std::ptr::NonNull; use std::sync::Arc; @@ -143,15 +142,12 @@ impl PyCatalogList { "Schema with name {name} doesn't exist." )))?; - Python::attach(|py| { - match catalog - .as_any() - .downcast_ref::() - { + Python::attach( + |py| match catalog.downcast_ref::() { Some(wrapped_catalog) => Ok(wrapped_catalog.catalog_provider.clone_ref(py)), None => PyCatalog::new_from_parts(catalog, self.codec.clone()).into_py_any(py), - } - }) + }, + ) } pub fn register_catalog(&self, name: &str, catalog_provider: Bound<'_, PyAny>) -> PyResult<()> { @@ -201,15 +197,12 @@ impl PyCatalog { "Schema with name {name} doesn't exist." )))?; - Python::attach(|py| { - match schema - .as_any() - .downcast_ref::() - { + Python::attach( + |py| match schema.downcast_ref::() { Some(wrapped_schema) => Ok(wrapped_schema.schema_provider.clone_ref(py)), None => PySchema::new_from_parts(schema, self.codec.clone()).into_py_any(py), - } - }) + }, + ) } pub fn register_schema(&self, name: &str, schema_provider: Bound<'_, PyAny>) -> PyResult<()> { @@ -356,10 +349,6 @@ impl SchemaProvider for RustWrappedPySchemaProvider { self.owner_name.as_deref() } - fn as_any(&self) -> &dyn Any { - self - } - fn table_names(&self) -> Vec { Python::attach(|py| { let provider = self.schema_provider.bind(py); @@ -465,10 +454,6 @@ impl RustWrappedPyCatalogProvider { #[async_trait] impl CatalogProvider for RustWrappedPyCatalogProvider { - fn as_any(&self) -> &dyn Any { - self - } - fn schema_names(&self) -> Vec { Python::attach(|py| { let provider = self.catalog_provider.bind(py); @@ -496,10 +481,7 @@ impl CatalogProvider for RustWrappedPyCatalogProvider { schema: Arc, ) -> datafusion::common::Result>> { Python::attach(|py| { - let py_schema = match schema - .as_any() - .downcast_ref::() - { + let py_schema = match schema.downcast_ref::() { Some(wrapped_schema) => wrapped_schema.schema_provider.as_any(), None => &PySchema::new_from_parts(schema, self.codec.clone()) .into_py_any(py) @@ -573,10 +555,6 @@ impl RustWrappedPyCatalogProviderList { #[async_trait] impl CatalogProviderList for RustWrappedPyCatalogProviderList { - fn as_any(&self) -> &dyn Any { - self - } - fn catalog_names(&self) -> Vec { Python::attach(|py| { let provider = self.catalog_provider_list.bind(py); @@ -604,10 +582,7 @@ impl CatalogProviderList for RustWrappedPyCatalogProviderList { catalog: Arc, ) -> Option> { Python::attach(|py| { - let py_catalog = match catalog - .as_any() - .downcast_ref::() - { + let py_catalog = match catalog.downcast_ref::() { Some(wrapped_schema) => wrapped_schema.catalog_provider.as_any().clone_ref(py), None => { match PyCatalog::new_from_parts(catalog, self.codec.clone()).into_py_any(py) { @@ -661,8 +636,8 @@ fn extract_catalog_provider_from_pyobj( .pointer_checked(Some(c"datafusion_catalog_provider"))? .cast(); let provider = unsafe { data.as_ref() }; - let provider: Arc = provider.into(); - provider as Arc + let provider: Arc = provider.into(); + provider } else { match catalog_provider.extract::() { Ok(py_catalog) => py_catalog.catalog, @@ -693,8 +668,8 @@ fn extract_schema_provider_from_pyobj( .pointer_checked(Some(c"datafusion_schema_provider"))? .cast(); let provider = unsafe { data.as_ref() }; - let provider: Arc = provider.into(); - provider as Arc + let provider: Arc = provider.into(); + provider } else { match schema_provider.extract::() { Ok(py_schema) => py_schema.schema, diff --git a/crates/core/src/common/data_type.rs b/crates/core/src/common/data_type.rs index af4179806..e79aea4ef 100644 --- a/crates/core/src/common/data_type.rs +++ b/crates/core/src/common/data_type.rs @@ -334,6 +334,9 @@ impl DataTypeMap { Ok(DataType::Interval(IntervalUnit::MonthDayNano)) } ScalarValue::List(arr) => Ok(arr.data_type().to_owned()), + ScalarValue::LargeList(arr) => Ok(arr.data_type().to_owned()), + ScalarValue::ListView(arr) => Ok(arr.data_type().to_owned()), + ScalarValue::LargeListView(arr) => Ok(arr.data_type().to_owned()), ScalarValue::Struct(_fields) => Err(PyNotImplementedError::new_err( "ScalarValue::Struct".to_string(), )), @@ -346,9 +349,6 @@ impl DataTypeMap { "ScalarValue::FixedSizeList".to_string(), )) } - ScalarValue::LargeList(_) => Err(PyNotImplementedError::new_err( - "ScalarValue::LargeList".to_string(), - )), ScalarValue::DurationSecond(_) => Ok(DataType::Duration(TimeUnit::Second)), ScalarValue::DurationMillisecond(_) => Ok(DataType::Duration(TimeUnit::Millisecond)), ScalarValue::DurationMicrosecond(_) => Ok(DataType::Duration(TimeUnit::Microsecond)), diff --git a/crates/core/src/common/schema.rs b/crates/core/src/common/schema.rs index 29a27b204..94b3ce0ae 100644 --- a/crates/core/src/common/schema.rs +++ b/crates/core/src/common/schema.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::borrow::Cow; use std::fmt::{self, Display, Formatter}; use std::sync::Arc; @@ -219,10 +218,6 @@ impl SqlTableSource { /// Implement TableSource, used in the logical query plan and in logical query optimizations impl TableSource for SqlTableSource { - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { self.schema.clone() } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 96de01889..642afeef7 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -773,8 +773,8 @@ impl PySessionContext { .pointer_checked(Some(c"datafusion_catalog_provider_list"))? .cast(); let provider = unsafe { data.as_ref() }; - let provider: Arc = provider.into(); - provider as Arc + let provider: Arc = provider.into(); + provider } else { match provider.extract::() { Ok(py_catalog_list) => py_catalog_list.catalog_list, @@ -809,8 +809,8 @@ impl PySessionContext { .pointer_checked(Some(c"datafusion_catalog_provider"))? .cast(); let provider = unsafe { data.as_ref() }; - let provider: Arc = provider.into(); - provider as Arc + let provider: Arc = provider.into(); + provider } else { match provider.extract::() { Ok(py_catalog) => py_catalog.catalog, @@ -1071,10 +1071,7 @@ impl PySessionContext { "Catalog with name {name} doesn't exist." )))?; - match catalog - .as_any() - .downcast_ref::() - { + match catalog.downcast_ref::() { Some(wrapped_schema) => Ok(wrapped_schema.catalog_provider.clone_ref(py)), None => { Ok(PyCatalog::new_from_parts(catalog, self.ffi_logical_codec()).into_py_any(py)?) diff --git a/crates/core/src/dataframe.rs b/crates/core/src/dataframe.rs index 2e74991b8..8f1a20d0d 100644 --- a/crates/core/src/dataframe.rs +++ b/crates/core/src/dataframe.rs @@ -578,6 +578,11 @@ impl PyDataFrame { Ok(PyTable::from(table_provider)) } + fn alias(&self, alias: &str) -> PyDataFusionResult { + let df = self.df.as_ref().clone().alias(alias)?; + Ok(Self::new(df)) + } + #[pyo3(signature = (*args))] fn select_exprs(&self, args: Vec) -> PyDataFusionResult { let args = args.iter().map(|s| s.as_ref()).collect::>(); diff --git a/crates/core/src/dataset.rs b/crates/core/src/dataset.rs index dbeafcd9f..2a5770338 100644 --- a/crates/core/src/dataset.rs +++ b/crates/core/src/dataset.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use async_trait::async_trait; @@ -62,12 +61,6 @@ impl Dataset { #[async_trait] impl TableProvider for Dataset { - /// Returns the table provider as [`Any`](std::any::Any) so that it can be - /// downcast to a specific implementation. - fn as_any(&self) -> &dyn Any { - self - } - /// Get a reference to the schema for this table fn schema(&self) -> SchemaRef { Python::attach(|py| { diff --git a/crates/core/src/dataset_exec.rs b/crates/core/src/dataset_exec.rs index a7dd1500d..771119a0f 100644 --- a/crates/core/src/dataset_exec.rs +++ b/crates/core/src/dataset_exec.rs @@ -15,18 +15,18 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::error::{ArrowError, Result as ArrowResult}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::{DataFusionError as InnerDataFusionError, Result as DFResult}; use datafusion::execution::context::TaskContext; use datafusion::logical_expr::Expr; use datafusion::logical_expr::utils::conjunction; -use datafusion::physical_expr::{EquivalenceProperties, LexOrdering}; +use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ @@ -156,11 +156,6 @@ impl ExecutionPlan for DatasetExec { Self::static_name() } - /// Return a reference to Any that can be used for downcasting - fn as_any(&self) -> &dyn Any { - self - } - /// Get the schema for this execution plan fn schema(&self) -> SchemaRef { self.schema.clone() @@ -235,8 +230,15 @@ impl ExecutionPlan for DatasetExec { }) } - fn partition_statistics(&self, _partition: Option) -> DFResult { - Ok(self.projected_statistics.clone()) + fn partition_statistics(&self, _partition: Option) -> DFResult> { + Ok(Arc::new(self.projected_statistics.clone())) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> DFResult, + ) -> DFResult { + Ok(TreeNodeRecursion::Continue) } fn properties(&self) -> &Arc { diff --git a/crates/core/src/expr.rs b/crates/core/src/expr.rs index 2e633baeb..eac571a11 100644 --- a/crates/core/src/expr.rs +++ b/crates/core/src/expr.rs @@ -23,8 +23,8 @@ use datafusion::arrow::datatypes::{DataType, Field}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::functions::core::expr_ext::FieldAccessor; use datafusion::logical_expr::expr::{ - AggregateFunction, AggregateFunctionParams, FieldMetadata, InList, InSubquery, ScalarFunction, - SetComparison, WindowFunction, + AggregateFunction, AggregateFunctionParams, FieldMetadata, HigherOrderFunction, InList, + InSubquery, Lambda, ScalarFunction, SetComparison, WindowFunction, }; use datafusion::logical_expr::utils::exprlist_to_fields; use datafusion::logical_expr::{ @@ -91,9 +91,12 @@ pub mod explain; pub mod extension; pub mod filter; pub mod grouping_set; +pub mod higher_order_function; pub mod in_list; pub mod in_subquery; pub mod join; +pub mod lambda; +pub mod lambda_variable; pub mod like; pub mod limit; pub mod literal; @@ -226,6 +229,14 @@ impl PyExpr { Expr::SetComparison(value) => { Ok(set_comparison::PySetComparison::from(value.clone()).into_bound_py_any(py)?) } + Expr::HigherOrderFunction(value) => Ok( + higher_order_function::PyHigherOrderFunction::from(value.clone()) + .into_bound_py_any(py)?, + ), + Expr::Lambda(value) => Ok(lambda::PyLambda::from(value.clone()).into_bound_py_any(py)?), + Expr::LambdaVariable(value) => { + Ok(lambda_variable::PyLambdaVariable::from(value.clone()).into_bound_py_any(py)?) + } }) } @@ -393,7 +404,10 @@ impl PyExpr { | Expr::OuterReferenceColumn(_, _) | Expr::Unnest(_) | Expr::IsNotUnknown(_) - | Expr::SetComparison(_) => RexType::Call, + | Expr::SetComparison(_) + | Expr::HigherOrderFunction(..) + | Expr::Lambda(..) => RexType::Call, + Expr::LambdaVariable(..) => RexType::Reference, Expr::ScalarSubquery(..) => RexType::ScalarSubquery, #[allow(deprecated)] Expr::Wildcard { .. } => { @@ -425,9 +439,10 @@ impl PyExpr { pub fn rex_call_operands(&self) -> PyResult> { match &self.expr { // Expr variants that are themselves the operand to return - Expr::Column(..) | Expr::ScalarVariable(..) | Expr::Literal(..) => { - Ok(vec![PyExpr::from(self.expr.clone())]) - } + Expr::Column(..) + | Expr::ScalarVariable(..) + | Expr::Literal(..) + | Expr::LambdaVariable(..) => Ok(vec![PyExpr::from(self.expr.clone())]), Expr::Alias(alias) => Ok(vec![PyExpr::from(*alias.expr.clone())]), @@ -454,13 +469,15 @@ impl PyExpr { params: AggregateFunctionParams { args, .. }, .. }) - | Expr::ScalarFunction(ScalarFunction { args, .. }) => { + | Expr::ScalarFunction(ScalarFunction { args, .. }) + | Expr::HigherOrderFunction(HigherOrderFunction { args, .. }) => { Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect()) } Expr::WindowFunction(boxed_window_fn) => { let args = &boxed_window_fn.params.args; Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect()) } + Expr::Lambda(Lambda { body, .. }) => Ok(vec![PyExpr::from(*body.clone())]), // Expr(s) that require more specific processing Expr::Case(Case { @@ -550,6 +567,10 @@ impl PyExpr { right: _, }) => format!("{op}"), Expr::ScalarFunction(ScalarFunction { func, args: _ }) => func.name().to_string(), + Expr::HigherOrderFunction(HigherOrderFunction { func, args: _ }) => { + func.name().to_string() + } + Expr::Lambda(..) => "lambda".to_string(), Expr::Cast { .. } => "cast".to_string(), Expr::Between { .. } => "between".to_string(), Expr::Case { .. } => "case".to_string(), @@ -837,7 +858,9 @@ impl PyExpr { | Operator::QuestionPipe | Operator::Colon => Err(py_type_err(format!("Unsupported expr: ${op}"))), }, - Expr::Cast(Cast { expr: _, data_type }) => DataTypeMap::map_from_arrow_type(data_type), + Expr::Cast(Cast { expr: _, field }) => { + DataTypeMap::map_from_arrow_type(field.data_type()) + } Expr::Literal(scalar_value, _) => DataTypeMap::map_from_scalar_value(scalar_value), _ => Err(py_type_err(format!( "Non Expr::Literal encountered in types: {expr:?}" @@ -893,6 +916,9 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/core/src/expr/cast.rs b/crates/core/src/expr/cast.rs index 37d603538..484d0c059 100644 --- a/crates/core/src/expr/cast.rs +++ b/crates/core/src/expr/cast.rs @@ -52,7 +52,7 @@ impl PyCast { } fn data_type(&self) -> PyResult { - Ok(self.cast.data_type.clone().into()) + Ok(self.cast.field.data_type().clone().into()) } } @@ -81,6 +81,6 @@ impl PyTryCast { } fn data_type(&self) -> PyResult { - Ok(self.try_cast.data_type.clone().into()) + Ok(self.try_cast.field.data_type().clone().into()) } } diff --git a/crates/core/src/expr/higher_order_function.rs b/crates/core/src/expr/higher_order_function.rs new file mode 100644 index 000000000..91a94de2c --- /dev/null +++ b/crates/core/src/expr/higher_order_function.rs @@ -0,0 +1,84 @@ +// 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. + +use std::fmt::{self, Display, Formatter}; + +use datafusion::logical_expr::expr::HigherOrderFunction; +use pyo3::prelude::*; + +use super::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "HigherOrderFunction", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyHigherOrderFunction { + higher_order: HigherOrderFunction, +} + +impl From for PyHigherOrderFunction { + fn from(higher_order: HigherOrderFunction) -> PyHigherOrderFunction { + PyHigherOrderFunction { higher_order } + } +} + +impl From for HigherOrderFunction { + fn from(higher_order: PyHigherOrderFunction) -> Self { + higher_order.higher_order + } +} + +impl Display for PyHigherOrderFunction { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!( + f, + "HigherOrderFunction(name={}, args={:?})", + self.higher_order.name(), + &self.higher_order.args, + ) + } +} + +#[pymethods] +impl PyHigherOrderFunction { + /// Name of the higher-order function being invoked. + fn name(&self) -> String { + self.higher_order.name().to_string() + } + + /// Arguments passed to the higher-order function. Some entries may be + /// `Lambda` expressions; others are ordinary value expressions. + fn args(&self) -> Vec { + self.higher_order + .args + .iter() + .map(|e| PyExpr::from(e.clone())) + .collect() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("HigherOrderFunction({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("HigherOrderFunction".to_string()) + } +} diff --git a/crates/core/src/expr/lambda.rs b/crates/core/src/expr/lambda.rs new file mode 100644 index 000000000..3ebc6e61c --- /dev/null +++ b/crates/core/src/expr/lambda.rs @@ -0,0 +1,78 @@ +// 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. + +use std::fmt::{self, Display, Formatter}; + +use datafusion::logical_expr::expr::Lambda; +use pyo3::prelude::*; + +use super::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "Lambda", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyLambda { + lambda: Lambda, +} + +impl From for PyLambda { + fn from(lambda: Lambda) -> PyLambda { + PyLambda { lambda } + } +} + +impl From for Lambda { + fn from(lambda: PyLambda) -> Self { + lambda.lambda + } +} + +impl Display for PyLambda { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!( + f, + "Lambda(params={:?}, body={:?})", + &self.lambda.params, &self.lambda.body, + ) + } +} + +#[pymethods] +impl PyLambda { + /// Parameter names of the lambda. + fn params(&self) -> Vec { + self.lambda.params.clone() + } + + /// Body expression of the lambda. + fn body(&self) -> PyExpr { + (*self.lambda.body).clone().into() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("Lambda({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("Lambda".to_string()) + } +} diff --git a/crates/core/src/expr/lambda_variable.rs b/crates/core/src/expr/lambda_variable.rs new file mode 100644 index 000000000..2ef554e17 --- /dev/null +++ b/crates/core/src/expr/lambda_variable.rs @@ -0,0 +1,67 @@ +// 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. + +use std::fmt::{self, Display, Formatter}; + +use datafusion::logical_expr::expr::LambdaVariable; +use pyo3::prelude::*; + +#[pyclass( + from_py_object, + frozen, + name = "LambdaVariable", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyLambdaVariable { + variable: LambdaVariable, +} + +impl From for PyLambdaVariable { + fn from(variable: LambdaVariable) -> PyLambdaVariable { + PyLambdaVariable { variable } + } +} + +impl From for LambdaVariable { + fn from(variable: PyLambdaVariable) -> Self { + variable.variable + } +} + +impl Display for PyLambdaVariable { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "LambdaVariable({})", &self.variable.name) + } +} + +#[pymethods] +impl PyLambdaVariable { + /// Reference name of the lambda parameter. + fn name(&self) -> String { + self.variable.name.clone() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("LambdaVariable({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("LambdaVariable".to_string()) + } +} diff --git a/crates/core/src/table.rs b/crates/core/src/table.rs index 623349771..e0f0f0d13 100644 --- a/crates/core/src/table.rs +++ b/crates/core/src/table.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use arrow::datatypes::SchemaRef; @@ -150,10 +149,6 @@ impl TempViewTable { #[async_trait] impl TableProvider for TempViewTable { - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { Arc::new(self.df.schema().as_arrow().clone()) } diff --git a/crates/core/src/udf.rs b/crates/core/src/udf.rs index c0a39cb47..d48bc729c 100644 --- a/crates/core/src/udf.rs +++ b/crates/core/src/udf.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::hash::{Hash, Hasher}; use std::ptr::NonNull; use std::sync::Arc; @@ -94,10 +93,6 @@ impl Hash for PythonFunctionScalarUDF { } impl ScalarUDFImpl for PythonFunctionScalarUDF { - fn as_any(&self) -> &dyn Any { - self - } - fn name(&self) -> &str { &self.name } diff --git a/crates/core/src/udtf.rs b/crates/core/src/udtf.rs index 9371732dc..b3de25e52 100644 --- a/crates/core/src/udtf.rs +++ b/crates/core/src/udtf.rs @@ -18,7 +18,7 @@ use std::ptr::NonNull; use std::sync::Arc; -use datafusion::catalog::{TableFunctionImpl, TableProvider}; +use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider}; use datafusion::error::Result as DataFusionResult; use datafusion::logical_expr::Expr; use datafusion_ffi::udtf::FFI_TableFunction; @@ -93,7 +93,11 @@ impl PyTableFunction { #[pyo3(signature = (*args))] pub fn __call__(&self, args: Vec) -> PyResult { let args: Vec = args.iter().map(|e| e.expr.clone()).collect(); - let table_provider = self.call(&args).map_err(py_datafusion_err)?; + let global = PySessionContext::global_ctx()?; + let state = global.ctx.state(); + let table_provider = self + .call_with_args(TableFunctionArgs::new(&args, &state)) + .map_err(py_datafusion_err)?; Ok(PyTable::from(table_provider)) } @@ -125,10 +129,12 @@ fn call_python_table_function( } impl TableFunctionImpl for PyTableFunction { - fn call(&self, args: &[Expr]) -> DataFusionResult> { + fn call_with_args(&self, args: TableFunctionArgs) -> DataFusionResult> { match &self.inner { - PyTableFunctionInner::FFIFunction(func) => func.call(args), - PyTableFunctionInner::PythonFunction(obj) => call_python_table_function(obj, args), + PyTableFunctionInner::FFIFunction(func) => func.call_with_args(args), + PyTableFunctionInner::PythonFunction(obj) => { + call_python_table_function(obj, args.exprs()) + } } } } diff --git a/crates/core/src/udwf.rs b/crates/core/src/udwf.rs index 1d3608ada..40e6208c4 100644 --- a/crates/core/src/udwf.rs +++ b/crates/core/src/udwf.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::ops::Range; use std::ptr::NonNull; use std::sync::Arc; @@ -317,10 +316,6 @@ impl MultiColumnWindowUDF { } impl WindowUDFImpl for MultiColumnWindowUDF { - fn as_any(&self) -> &dyn Any { - self - } - fn name(&self) -> &str { &self.name } diff --git a/examples/datafusion-ffi-example/python/tests/_test_schema_provider.py b/examples/datafusion-ffi-example/python/tests/_test_schema_provider.py index 93449c660..c4a94348d 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_schema_provider.py +++ b/examples/datafusion-ffi-example/python/tests/_test_schema_provider.py @@ -63,15 +63,18 @@ def test_schema_provider_extract_values(inner_capsule: bool) -> None: result = ctx.table(f"{expected_schema_name}.{expected_table_name}").collect() assert len(result) == 2 + # Multi-partition collect order is non-deterministic; sort batches by + # first value of col0 so col0 and col1 stay aligned. + result = sorted(result, key=lambda r: r.column(0)[0].as_py()) col0_result = [r.column(0) for r in result] col1_result = [r.column(1) for r in result] expected_col0 = [ - pa.array([10, 20, 30], type=pa.int32()), pa.array([5, 7], type=pa.int32()), + pa.array([10, 20, 30], type=pa.int32()), ] expected_col1 = [ - pa.array([1, 2, 5], type=pa.float64()), pa.array([1.5, 2.5], type=pa.float64()), + pa.array([1, 2, 5], type=pa.float64()), ] assert col0_result == expected_col0 assert col1_result == expected_col1 diff --git a/examples/datafusion-ffi-example/python/tests/_test_table_function.py b/examples/datafusion-ffi-example/python/tests/_test_table_function.py index bf5aae3bd..de4662dd7 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_table_function.py +++ b/examples/datafusion-ffi-example/python/tests/_test_table_function.py @@ -39,7 +39,8 @@ def test_ffi_table_function_register() -> None: assert result[0].num_columns == 4 print(result) - result = [r.column(0) for r in result] + # Multi-partition collect order is non-deterministic; sort by first value. + result = sorted((r.column(0) for r in result), key=lambda a: a[0].as_py()) expected = [ pa.array([0, 1, 2], type=pa.int32()), pa.array([3, 4, 5, 6], type=pa.int32()), @@ -61,7 +62,7 @@ def test_ffi_table_function_call_directly(): assert result[0].num_columns == 4 print(result) - result = [r.column(0) for r in result] + result = sorted((r.column(0) for r in result), key=lambda a: a[0].as_py()) expected = [ pa.array([0, 1, 2], type=pa.int32()), pa.array([3, 4, 5, 6], type=pa.int32()), @@ -96,7 +97,7 @@ def common_table_function_test(test_ctx: SessionContext) -> None: assert result[0].num_columns == 3 print(result) - result = [r.column(0) for r in result] + result = sorted((r.column(0) for r in result), key=lambda a: a[0].as_py()) expected = [ pa.array([0, 1], type=pa.int32()), pa.array([2, 3, 4], type=pa.int32()), diff --git a/examples/datafusion-ffi-example/python/tests/_test_table_provider.py b/examples/datafusion-ffi-example/python/tests/_test_table_provider.py index fc77d2d3b..aee16f839 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_table_provider.py +++ b/examples/datafusion-ffi-example/python/tests/_test_table_provider.py @@ -36,7 +36,9 @@ def test_table_provider_ffi(inner_capsule: bool) -> None: assert len(result) == 4 assert result[0].num_columns == 3 - result = [r.column(0) for r in result] + # Multi-partition collect order is non-deterministic; sort by first value + # in column 0, which is unique per partition (0, 2, 4, 6). + result = sorted((r.column(0) for r in result), key=lambda a: a[0].as_py()) expected = [ pa.array([0, 1], type=pa.int32()), pa.array([2, 3, 4], type=pa.int32()), @@ -47,5 +49,5 @@ def test_table_provider_ffi(inner_capsule: bool) -> None: assert result == expected result = ctx.read_table(table).collect() - result = [r.column(0) for r in result] + result = sorted((r.column(0) for r in result), key=lambda a: a[0].as_py()) assert result == expected diff --git a/examples/datafusion-ffi-example/src/aggregate_udf.rs b/examples/datafusion-ffi-example/src/aggregate_udf.rs index d5343ff91..86737778f 100644 --- a/examples/datafusion-ffi-example/src/aggregate_udf.rs +++ b/examples/datafusion-ffi-example/src/aggregate_udf.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use arrow_schema::DataType; @@ -61,10 +60,6 @@ impl MySumUDF { } impl AggregateUDFImpl for MySumUDF { - fn as_any(&self) -> &dyn Any { - self - } - fn name(&self) -> &str { "my_custom_sum" } diff --git a/examples/datafusion-ffi-example/src/catalog_provider.rs b/examples/datafusion-ffi-example/src/catalog_provider.rs index bd5da1e4d..6131ab0f0 100644 --- a/examples/datafusion-ffi-example/src/catalog_provider.rs +++ b/examples/datafusion-ffi-example/src/catalog_provider.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::fmt::Debug; use std::sync::Arc; @@ -106,10 +105,6 @@ impl FixedSchemaProvider { #[async_trait] impl SchemaProvider for FixedSchemaProvider { - fn as_any(&self) -> &dyn Any { - self - } - fn table_names(&self) -> Vec { self.inner.table_names() } @@ -149,10 +144,6 @@ pub(crate) struct MyCatalogProvider { } impl CatalogProvider for MyCatalogProvider { - fn as_any(&self) -> &dyn Any { - self - } - fn schema_names(&self) -> Vec { self.inner.schema_names() } @@ -220,10 +211,6 @@ pub(crate) struct MyCatalogProviderList { } impl CatalogProviderList for MyCatalogProviderList { - fn as_any(&self) -> &dyn Any { - self - } - fn catalog_names(&self) -> Vec { self.inner.catalog_names() } diff --git a/examples/datafusion-ffi-example/src/scalar_udf.rs b/examples/datafusion-ffi-example/src/scalar_udf.rs index 374924781..a3c65e875 100644 --- a/examples/datafusion-ffi-example/src/scalar_udf.rs +++ b/examples/datafusion-ffi-example/src/scalar_udf.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use arrow_array::{Array, BooleanArray}; @@ -61,10 +60,6 @@ impl IsNullUDF { } impl ScalarUDFImpl for IsNullUDF { - fn as_any(&self) -> &dyn Any { - self - } - fn name(&self) -> &str { "my_custom_is_null" } diff --git a/examples/datafusion-ffi-example/src/window_udf.rs b/examples/datafusion-ffi-example/src/window_udf.rs index cbf179a86..f33a166ed 100644 --- a/examples/datafusion-ffi-example/src/window_udf.rs +++ b/examples/datafusion-ffi-example/src/window_udf.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use arrow_schema::{DataType, FieldRef}; @@ -56,10 +55,6 @@ impl MyRankUDF { } impl WindowUDFImpl for MyRankUDF { - fn as_any(&self) -> &dyn Any { - self - } - fn name(&self) -> &str { "my_custom_rank" } diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index f08b464bb..601419fab 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -69,7 +69,12 @@ # The following imports are okay to remain as opaque to the user. from ._internal import Config -from .catalog import Catalog, Table +from .catalog import ( + Catalog, + Table, + TableProviderFactory, + TableProviderFactoryExportable, +) from .col import col, column from .common import DFSchema from .context import ( @@ -133,6 +138,8 @@ "SessionContext", "Table", "TableFunction", + "TableProviderFactory", + "TableProviderFactoryExportable", "WindowFrame", "WindowUDF", "catalog", diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index 2b07861da..9ac8293d6 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -523,6 +523,32 @@ def select_exprs(self, *args: str) -> DataFrame: """ return self.df.select_exprs(*args) + def alias(self, alias: str) -> DataFrame: + """Assign a table alias to this :py:class:`DataFrame`. + + Replaces the qualifiers of the output columns with ``alias``. Useful for + self-joins and any situation that needs an unambiguous table-style + qualifier (``alias.col``) for downstream references. + + Args: + alias: Table alias to apply to the DataFrame's columns. + + Returns: + DataFrame with columns re-qualified under ``alias``. + + Example: + >>> from datafusion import col + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"id": [1, 2], "val": [10, 20]}) + >>> left = df.alias("l") + >>> right = df.alias("r") + >>> left.join(right, left_on="id", right_on="id").select( + ... "id", col("l.val").alias("lval"), col("r.val").alias("rval") + ... ).sort("id").to_pydict() + {'id': [1, 2], 'lval': [10, 20], 'rval': [10, 20]} + """ + return DataFrame(self.df.alias(alias)) + def select(self, *exprs: Expr | str) -> DataFrame: """Project arbitrary expressions into a new :py:class:`DataFrame`. diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index e0135e3ed..55cc2e52a 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -150,6 +150,9 @@ TransactionStart = expr_internal.TransactionStart TryCast = expr_internal.TryCast Union = expr_internal.Union +HigherOrderFunction = expr_internal.HigherOrderFunction +Lambda = expr_internal.Lambda +LambdaVariable = expr_internal.LambdaVariable Unnest = expr_internal.Unnest UnnestExpr = expr_internal.UnnestExpr Values = expr_internal.Values @@ -193,6 +196,7 @@ "FileType", "Filter", "GroupingSet", + "HigherOrderFunction", "ILike", "InList", "InSubquery", @@ -207,6 +211,8 @@ "Join", "JoinConstraint", "JoinType", + "Lambda", + "LambdaVariable", "Like", "Limit", "Literal", diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 08062851a..9761d1879 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -184,6 +184,7 @@ "ifnull", "in_list", "initcap", + "instr", "isnan", "iszero", "lag", @@ -273,6 +274,7 @@ "percent_rank", "percentile_cont", "pi", + "position", "pow", "power", "quantile_cont", @@ -3874,7 +3876,7 @@ def arrays_zip(*arrays: Expr) -> Expr: >>> result = df.select( ... dfn.functions.arrays_zip(dfn.col("a"), dfn.col("b")).alias("result")) >>> result.collect_column("result")[0].as_py() - [{'c0': 1, 'c1': 3}, {'c0': 2, 'c1': 4}] + [{'1': 1, '2': 3}, {'1': 2, '2': 4}] """ args = [a.expr for a in arrays] return Expr(f.arrays_zip(args)) diff --git a/python/tests/test_aggregation.py b/python/tests/test_aggregation.py index 240332848..f5c54f756 100644 --- a/python/tests/test_aggregation.py +++ b/python/tests/test_aggregation.py @@ -125,18 +125,34 @@ def test_aggregation_stats(df, agg_expr, calc_expected): pa.array([1], type=pa.uint64()), False, ), - (f.approx_median(column("b")), pa.array([4]), False), - (f.median(column("b"), distinct=True), pa.array([5]), False), - (f.median(column("b"), filter=column("a") != 2), pa.array([5]), False), - (f.approx_median(column("b"), filter=column("a") != 2), pa.array([5]), False), - (f.approx_percentile_cont(column("b"), 0.5), pa.array([4]), False), + (f.approx_median(column("b")), pa.array([4], type=pa.float64()), False), + ( + f.median(column("b"), distinct=True), + pa.array([5], type=pa.float64()), + False, + ), + ( + f.median(column("b"), filter=column("a") != 2), + pa.array([5], type=pa.float64()), + False, + ), + ( + f.approx_median(column("b"), filter=column("a") != 2), + pa.array([5], type=pa.float64()), + False, + ), + ( + f.approx_percentile_cont(column("b"), 0.5), + pa.array([4], type=pa.float64()), + False, + ), ( f.approx_percentile_cont( column("b").sort(ascending=True, nulls_first=False), 0.5, num_centroids=2, ), - pa.array([4]), + pa.array([4.75], type=pa.float64()), False, ), ( @@ -212,19 +228,19 @@ def test_aggregation(df, agg_expr, expected, array_sort): ( "approx_percentile_cont", f.approx_percentile_cont(column("c3"), 0.95, num_centroids=200), - [73, 68, 122, 124, 115], + [73.55, 68.0, 122.5, 124.2, 115.6], ), ( "approx_perc_cont_few_centroids", f.approx_percentile_cont(column("c3"), 0.95, num_centroids=5), - [72, 68, 119, 124, 115], + [72.775, 68.0, 119.4075, 124.825, 115.44], ), ( "approx_perc_cont_filtered", f.approx_percentile_cont( column("c3"), 0.95, num_centroids=200, filter=column("c3") > lit(0) ), - [83, 68, 122, 124, 117], + [83.0, 68.0, 122.75, 124.9, 117.6], ), ( "corr", diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index 9e2f791ea..6bd0ce9f9 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -292,6 +292,20 @@ def test_select_exprs(df): assert result.column(1) == pa.array([3, 3, 3]) +def test_alias_self_join(df): + left = df.alias("l") + right = df.alias("r") + joined = left.join(right, left_on="a", right_on="a").select( + "a", + column("l.b").alias("lb"), + column("r.b").alias("rb"), + ) + result = joined.sort(column("a")).collect()[0] + assert result.column(0) == pa.array([1, 2, 3]) + assert result.column(1) == pa.array([4, 5, 6]) + assert result.column(2) == pa.array([4, 5, 6]) + + def test_drop_quoted_columns(): ctx = SessionContext() batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], names=["ID_For_Students"]) diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index 6a466f6f2..cf1fe43e8 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -172,7 +172,10 @@ def test_relational_expr(test_ctx): assert df.filter(col("b") == "beta").count() == 1 assert df.filter(col("b") != "beta").count() == 2 - assert df.filter(col("a") == "beta").count() == 0 + # Upstream DataFusion now errors on string→Int64 implicit cast in filter + # (previously silently produced 0 matches). + with pytest.raises(Exception, match="Cannot cast string 'beta'"): + df.filter(col("a") == "beta").count() assert df.filter(col("a") == None).count() == 1 # noqa: E711 assert df.filter(col("a") != None).count() == 3 # noqa: E711 assert df.filter(col("b") == None).count() == 1 # noqa: E711 @@ -613,7 +616,7 @@ def test_alias_with_metadata(df): # pytest.param( col("c").reverse(), - pa.array(["olleH", " dlrow ", "!", None], type=pa.string()), + pa.array(["olleH", " dlrow ", "!", None], type=pa.string_view()), id="reverse", ), pytest.param( @@ -633,7 +636,7 @@ def test_alias_with_metadata(df): ), pytest.param( col("c").lower(), - pa.array(["hello", " world ", "!", None], type=pa.string()), + pa.array(["hello", " world ", "!", None], type=pa.string_view()), id="lower", ), pytest.param( @@ -767,7 +770,7 @@ def test_alias_with_metadata(df): ), pytest.param( col("c").upper(), - pa.array(["HELLO", " WORLD ", "!", None], type=pa.string()), + pa.array(["HELLO", " WORLD ", "!", None], type=pa.string_view()), id="upper", ), pytest.param( diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index d9781b1fb..5538fc33b 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -836,7 +836,7 @@ def test_map_functions(func, expected): (f.chr(literal(68)), pa.array(["D", "D", "D"])), ( f.concat_ws("-", column("a"), literal("test")), - pa.array(["Hello-test", "World-test", "!-test"]), + pa.array(["Hello-test", "World-test", "!-test"], type=pa.string_view()), ), ( f.concat(column("a").cast(pa.string()), literal("?")), @@ -851,7 +851,10 @@ def test_map_functions(func, expected): pa.array(["Hel", "Wor", "!"], type=pa.string_view()), ), (f.length(column("c")), pa.array([6, 7, 2], type=pa.int32())), - (f.lower(column("a")), pa.array(["hello", "world", "!"])), + ( + f.lower(column("a")), + pa.array(["hello", "world", "!"], type=pa.string_view()), + ), (f.lpad(column("a"), literal(7)), pa.array([" Hello", " World", " !"])), ( f.ltrim(column("c")), @@ -871,13 +874,16 @@ def test_map_functions(func, expected): (f.octet_length(column("a")), pa.array([5, 5, 1], type=pa.int32())), ( f.repeat(column("a"), literal(2)), - pa.array(["HelloHello", "WorldWorld", "!!"]), + pa.array(["HelloHello", "WorldWorld", "!!"], type=pa.string_view()), ), ( f.replace(column("a"), literal("l"), literal("?")), pa.array(["He??o", "Wor?d", "!"]), ), - (f.reverse(column("a")), pa.array(["olleH", "dlroW", "!"])), + ( + f.reverse(column("a")), + pa.array(["olleH", "dlroW", "!"], type=pa.string_view()), + ), ( f.right(column("a"), literal(4)), pa.array(["ello", "orld", "!"], type=pa.string_view()), @@ -892,7 +898,7 @@ def test_map_functions(func, expected): ), ( f.split_part(column("a"), literal("l"), literal(1)), - pa.array(["He", "Wor", "!"]), + pa.array(["He", "Wor", "!"], type=pa.string_view()), ), (f.contains(column("a"), literal("ell")), pa.array([True, False, False])), (f.starts_with(column("a"), literal("Wor")), pa.array([False, True, False])), @@ -903,14 +909,17 @@ def test_map_functions(func, expected): ), ( f.translate(column("a"), literal("or"), literal("ld")), - pa.array(["Helll", "Wldld", "!"]), + pa.array(["Helll", "Wldld", "!"], type=pa.string_view()), ), (f.trim(column("c")), pa.array(["hello", "world", "!"], type=pa.string_view())), - (f.upper(column("c")), pa.array(["HELLO ", " WORLD ", " !"])), + ( + f.upper(column("c")), + pa.array(["HELLO ", " WORLD ", " !"], type=pa.string_view()), + ), (f.ends_with(column("a"), literal("llo")), pa.array([True, False, False])), ( f.overlay(column("a"), literal("--"), literal(2)), - pa.array(["H--lo", "W--ld", "--"]), + pa.array(["H--lo", "W--ld", "!--"]), ), ( f.regexp_like(column("a"), literal("(ell|orl)")), @@ -2063,7 +2072,7 @@ def test_arrays_zip_aliases(func): df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]}) result = df.select(func(column("a"), column("b")).alias("v")).collect() values = result[0].column(0)[0].as_py() - assert values == [{"c0": 1, "c1": 3}, {"c0": 2, "c1": 4}] + assert values == [{"1": 1, "2": 3}, {"1": 2, "2": 4}] @pytest.mark.parametrize("func", [f.string_to_array, f.string_to_list]) diff --git a/python/tests/test_imports.py b/python/tests/test_imports.py index fca94b35a..fea4cc91f 100644 --- a/python/tests/test_imports.py +++ b/python/tests/test_imports.py @@ -22,6 +22,8 @@ DataFrame, ScalarUDF, SessionContext, + TableProviderFactory, + TableProviderFactoryExportable, functions, ) from datafusion.common import ( @@ -47,6 +49,7 @@ Extension, Filter, GroupingSet, + HigherOrderFunction, ILike, InList, InSubquery, @@ -60,6 +63,8 @@ Join, JoinConstraint, JoinType, + Lambda, + LambdaVariable, Like, Limit, Literal, @@ -160,9 +165,16 @@ def test_class_module_is_datafusion(): DropTable, Repartition, Partitioning, + HigherOrderFunction, + Lambda, + LambdaVariable, ]: assert klass.__module__ == "datafusion.expr" + # catalog factory types (re-exported at top level) + for klass in [TableProviderFactory, TableProviderFactoryExportable]: + assert klass.__module__ == "datafusion.catalog" + # schema for klass in [DFSchema]: assert klass.__module__ == "datafusion.common" From afaeccbf27623e983bba594f466024d1f52c5a0f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 19 May 2026 09:31:30 -0400 Subject: [PATCH 36/83] feat: enable pickling of most Expr except udaf and udwf (#1544) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: pickle support for Expr via inline scalar UDF encoding Adds Python-aware encoding to PythonLogicalCodec/PythonPhysicalCodec so a ScalarUDF defined in Python travels inside the serialized expression (cloudpickled into fun_definition) instead of needing a matching registration on the receiver. With that in place, Expr gains __reduce__ + classmethod from_bytes(buf, ctx=None) so pickle.dumps / pickle.loads work end-to-end on expressions built from col, lit, built-in functions, and Python scalar UDFs. Wire format is framed as ; the version byte lets a too-new/too-old payload surface a clean Execution error instead of an opaque cloudpickle unpack failure. Schema serde is via arrow-rs's native IPC (no pyarrow round-trip). Cloudpickle module handle is cached per-interpreter through PyOnceLock. Worker-side context resolution lives in a new datafusion.ipc module: set_worker_ctx / get_worker_ctx / clear_worker_ctx plus a private _resolve_ctx helper consulted by Expr.from_bytes. Priority is explicit ctx > worker ctx > global SessionContext. FFI UDFs still travel by name and require the matching registration on the receiver's context. Aggregate and window UDF inline encoding, the per-session with_python_udf_inlining toggle, sender-side context, and the user-guide docs land in follow-on PRs. Co-Authored-By: Claude Opus 4.7 (1M context) * docs(pickle): add cloudpickle security warnings, docstring examples, edge-case tests Inline `.. warning::` blocks on `Expr.to_bytes`, `Expr.from_bytes`, and `Expr.__reduce__` so the cloudpickle / arbitrary-code-execution caveat is visible at the public API surface in advance of the user-guide page that lands in PR 4. Add doctest-style `Examples:` blocks to `datafusion.ipc` functions (`set_worker_ctx`, `clear_worker_ctx`, `get_worker_ctx`, `_resolve_ctx`), `ScalarUDF.name`, and the new `Expr` pickle methods, per CLAUDE.md. Tighten `Expr.__reduce__` return annotation to `tuple[Callable[[bytes], Expr], tuple[bytes]]`. Tests: multi-arg UDF round-trip (covers synthetic `arg_{i}` schema-field loop in the codec) plus malformed-bytes paths through `Expr.from_bytes`. Co-Authored-By: Claude Opus 4.7 (1M context) * as_any no longer in api * feat(pickle): stamp Python (major, minor) in UDF wire header cloudpickle bytecode is not portable across Python minor versions — a payload produced on 3.11 fails to load on 3.12 with an opaque marshal/unpickle error. Embed the sender's (major, minor) in the DFPYUDF wire header and reject mismatches at decode time with an actionable error that names both versions, instead of letting the failure surface from inside cloudpickle.loads. Header layout becomes: DFPYUDF (7) | version (1) | py_major (1) | py_minor (1) | cloudpickle Extend the Security warnings on Expr.to_bytes / from_bytes / __reduce__ with a Portability section covering the cross-version constraint and cloudpickle's by-value/by-reference behavior (the callable inlines bytecode and closure cells, but imported names travel by reference and must be importable on the receiver). Add a matching Serialization model note to the datafusion.ipc module docstring. New tests: - codec::wire_header_tests: py-major/minor mismatch, truncated py-version bytes, round-trip with py-version - test_pickle_expr::test_cross_version_error_message: patches the py_minor byte inside an emitted payload and asserts the error message identifies the version mismatch Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/codec.rs | 524 ++++++++++++++++++++++++++++-- crates/core/src/udf.rs | 82 ++++- pyproject.toml | 7 + python/datafusion/__init__.py | 3 +- python/datafusion/expr.py | 164 +++++++++- python/datafusion/ipc.py | 161 +++++++++ python/datafusion/user_defined.py | 23 ++ python/tests/test_expr.py | 4 +- python/tests/test_pickle_expr.py | 184 +++++++++++ uv.lock | 13 +- 10 files changed, 1108 insertions(+), 57 deletions(-) create mode 100644 python/datafusion/ipc.py create mode 100644 python/tests/test_pickle_expr.py diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 088532df2..cc038edc9 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -19,11 +19,11 @@ //! //! Datafusion-python plans can carry references to Python-defined //! objects that the upstream protobuf codecs do not know how to -//! serialize: pure-Python scalar / aggregate / window UDFs, Python -//! query-planning extensions, and so on. Their state lives inside -//! `Py` callables and closures rather than being recoverable -//! from a name in the receiver's function registry. To ship a plan -//! across a process boundary (pickle, `multiprocessing`, Ray actor, +//! serialize: pure-Python scalar UDFs, Python query-planning +//! extensions, and so on. Their state lives inside `Py` +//! callables and closures rather than being recoverable from a name +//! in the receiver's function registry. To ship a plan across a +//! process boundary (pickle, `multiprocessing`, Ray actor, //! `datafusion-distributed`, etc.) those payloads have to be encoded //! into the proto wire format itself. //! @@ -48,52 +48,160 @@ //! plans to survive a serialization round-trip. Both codecs share //! the same payload framing for that reason. //! -//! Payloads emitted by these codecs are tagged with an 8-byte magic -//! prefix so the decoder can distinguish them from arbitrary bytes -//! (empty `fun_definition` from the default codec, user FFI payloads -//! that picked a non-colliding prefix). Dispatch precedence on -//! decode: **Python-inline payload (magic prefix match) → `inner` -//! codec → caller's `FunctionRegistry` fallback.** +//! Payloads emitted by these codecs are framed as +//! ` `. +//! The family magic identifies the UDF flavor; the version byte lets +//! the decoder reject too-new or too-old payloads with a clean error +//! instead of falling into an opaque `cloudpickle` tuple-unpack +//! failure when the tuple shape changes; the Python `(major, minor)` +//! bytes catch the cloudpickle-cross-minor-version case and raise an +//! actionable error instead of an opaque `marshal` failure on load +//! (cloudpickle payloads are not portable across Python minor +//! versions). Dispatch precedence on decode: **family match + +//! supported version + matching Python version → `inner` codec → +//! caller's `FunctionRegistry` fallback.** //! -//! ## Wire-format magic prefix registry +//! ## Wire-format family registry //! -//! | Layer + kind | Magic prefix | -//! | ----------------------------- | ------------ | -//! | `PythonLogicalCodec` scalar | `DFPYUDF1` | -//! | `PythonLogicalCodec` agg | `DFPYUDA1` | -//! | `PythonLogicalCodec` window | `DFPYUDW1` | -//! | `PythonPhysicalCodec` scalar | `DFPYUDF1` | -//! | `PythonPhysicalCodec` agg | `DFPYUDA1` | -//! | `PythonPhysicalCodec` window | `DFPYUDW1` | -//! | `PythonPhysicalCodec` expr | `DFPYPE1` | -//! | User FFI extension codec | user-chosen | -//! | Default codec | (none) | +//! | Layer + kind | Family prefix | +//! | ----------------------------- | ------------- | +//! | `PythonLogicalCodec` scalar | `DFPYUDF` | +//! | `PythonPhysicalCodec` scalar | `DFPYUDF` | +//! | User FFI extension codec | user-chosen | +//! | Default codec | (none) | //! -//! Downstream FFI codecs should pick non-colliding prefixes (use a -//! `DF` namespace plus a crate-specific suffix). The codec +//! Aggregate and window UDF families are reserved for follow-on work. +//! +//! Current wire-format version is [`WIRE_VERSION_CURRENT`]; supported +//! receive range is `WIRE_VERSION_MIN_SUPPORTED..=WIRE_VERSION_CURRENT`. +//! Bump [`WIRE_VERSION_CURRENT`] whenever the cloudpickle tuple shape +//! changes; raise [`WIRE_VERSION_MIN_SUPPORTED`] when dropping support +//! for an older shape. +//! +//! Downstream FFI codecs should pick non-colliding family prefixes +//! (use a `DF` namespace plus a crate-specific suffix). The codec //! implementations in this module currently delegate every method to //! `inner`; the encoder/decoder hooks for each kind are added as the //! corresponding Python-side type becomes serializable. use std::sync::Arc; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::ipc::reader::StreamReader; +use arrow::ipc::writer::StreamWriter; use datafusion::common::{Result, TableReference}; use datafusion::datasource::TableProvider; use datafusion::datasource::file_format::FileFormatFactory; use datafusion::execution::TaskContext; -use datafusion::logical_expr::{AggregateUDF, Extension, LogicalPlan, ScalarUDF, WindowUDF}; +use datafusion::logical_expr::{ + AggregateUDF, Extension, LogicalPlan, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, + Volatility, WindowUDF, +}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; use datafusion_proto::physical_plan::{DefaultPhysicalExtensionCodec, PhysicalExtensionCodec}; +use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; +use pyo3::types::{PyBytes, PyTuple}; + +use crate::udf::PythonFunctionScalarUDF; -/// Wire-format prefix that tags a `fun_definition` payload as an -/// inlined Python scalar UDF (cloudpickled tuple of name, callable, -/// input schema, return field, volatility). Defined once here so -/// the encoder and decoder cannot drift. -#[allow(dead_code)] -pub(crate) const PY_SCALAR_UDF_MAGIC: &[u8] = b"DFPYUDF1"; +// Wire-format framing for inlined Python UDF payloads. +// +// Layout: ` `. +// The family magic identifies the UDF flavor; the version byte lets +// the decoder reject too-new or too-old payloads with a clean error +// instead of falling into an opaque `cloudpickle` tuple-unpack failure +// when the tuple shape changes; the Python `(major, minor)` bytes +// catch the cloudpickle-cross-minor-version case (cloudpickle is not +// portable across Python minor versions) and raise an actionable +// error instead of an opaque `marshal` failure on load. Bump +// [`WIRE_VERSION_CURRENT`] whenever the tuple shape changes; raise +// [`WIRE_VERSION_MIN_SUPPORTED`] when dropping support for an older +// shape. + +/// Family prefix for an inlined Python scalar UDF +/// (cloudpickled tuple of name, callable, input schema, return field, +/// volatility). +pub(crate) const PY_SCALAR_UDF_FAMILY: &[u8] = b"DFPYUDF"; + +/// Wire-format version this build emits. +pub(crate) const WIRE_VERSION_CURRENT: u8 = 1; + +/// Oldest wire-format version this build still decodes. Bump when +/// retiring support for an older payload shape. +pub(crate) const WIRE_VERSION_MIN_SUPPORTED: u8 = 1; + +/// Tag `buf` with the framing header for `family` at the current +/// wire-format version, stamping `py_version` as `(major, minor)` +/// bytes. Append-only — the caller writes the cloudpickle payload +/// after. +fn write_wire_header(buf: &mut Vec, family: &[u8], py_version: (u8, u8)) { + buf.extend_from_slice(family); + buf.push(WIRE_VERSION_CURRENT); + buf.push(py_version.0); + buf.push(py_version.1); +} + +/// Inspect the framing on `buf`. +/// +/// * `Ok(None)` — `buf` does not carry `family`. The caller should +/// delegate to its `inner` codec. +/// * `Ok(Some(payload))` — `buf` carries `family` at a version this +/// build accepts and a Python `(major, minor)` matching +/// `expected_py`; `payload` is the cloudpickle blob. +/// * `Err(_)` — `buf` carries `family` but the wire-format version +/// is outside `WIRE_VERSION_MIN_SUPPORTED..=WIRE_VERSION_CURRENT`, +/// or the stamped Python `(major, minor)` does not match +/// `expected_py`. The error names the offending values so an +/// operator can diagnose sender/receiver drift instead of seeing +/// an opaque cloudpickle tuple-unpack or `marshal` failure. +fn strip_wire_header<'a>( + buf: &'a [u8], + family: &[u8], + kind: &str, + expected_py: (u8, u8), +) -> Result> { + if !buf.starts_with(family) { + return Ok(None); + } + let version_idx = family.len(); + let Some(&version) = buf.get(version_idx) else { + return Err(datafusion::error::DataFusionError::Execution(format!( + "Truncated inline Python {kind} payload: missing wire-format version byte" + ))); + }; + if !(WIRE_VERSION_MIN_SUPPORTED..=WIRE_VERSION_CURRENT).contains(&version) { + return Err(datafusion::error::DataFusionError::Execution(format!( + "Inline Python {kind} payload wire-format version v{version}; \ + this build supports v{WIRE_VERSION_MIN_SUPPORTED}..=v{WIRE_VERSION_CURRENT}. \ + Align datafusion-python versions on sender and receiver." + ))); + } + let py_major_idx = version_idx + 1; + let Some(&encoded_major) = buf.get(py_major_idx) else { + return Err(datafusion::error::DataFusionError::Execution(format!( + "Truncated inline Python {kind} payload: missing Python major version byte" + ))); + }; + let py_minor_idx = version_idx + 2; + let Some(&encoded_minor) = buf.get(py_minor_idx) else { + return Err(datafusion::error::DataFusionError::Execution(format!( + "Truncated inline Python {kind} payload: missing Python minor version byte" + ))); + }; + let (current_major, current_minor) = expected_py; + if encoded_major != current_major || encoded_minor != current_minor { + return Err(datafusion::error::DataFusionError::Execution(format!( + "Inline Python {kind} payload was serialized on Python \ + {encoded_major}.{encoded_minor} but this process is running Python \ + {current_major}.{current_minor}. cloudpickle payloads are not portable \ + across Python minor versions. Align Python versions on sender and receiver." + ))); + } + Ok(Some(&buf[py_minor_idx + 1..])) +} /// `LogicalExtensionCodec` parked on every `SessionContext`. Holds /// the Python-aware encoding hooks for logical-layer types @@ -177,10 +285,16 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + if try_encode_python_scalar_udf(node, buf)? { + return Ok(()); + } self.inner.try_encode_udf(node, buf) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + if let Some(udf) = try_decode_python_scalar_udf(buf)? { + return Ok(udf); + } self.inner.try_decode_udf(name, buf) } @@ -212,7 +326,7 @@ impl LogicalExtensionCodec for PythonLogicalCodec { /// encoding on this layer too — otherwise a plan with a Python UDF /// would round-trip at the logical level but break at the physical /// level. Both layers reuse the shared payload framing -/// ([`PY_SCALAR_UDF_MAGIC`] et al.) so the wire format is identical. +/// ([`PY_SCALAR_UDF_FAMILY`]) so the wire format is identical. #[derive(Debug)] pub struct PythonPhysicalCodec { inner: Arc, @@ -249,10 +363,16 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + if try_encode_python_scalar_udf(node, buf)? { + return Ok(()); + } self.inner.try_encode_udf(node, buf) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + if let Some(udf) = try_decode_python_scalar_udf(buf)? { + return Ok(udf); + } self.inner.try_decode_udf(name, buf) } @@ -284,3 +404,339 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { self.inner.try_decode_udwf(name, buf) } } + +// ============================================================================= +// Shared Python scalar UDF encode / decode helpers +// +// Both `PythonLogicalCodec` and `PythonPhysicalCodec` consult these on +// every `try_encode_udf` / `try_decode_udf` call. Same wire format on +// both layers — a Python `ScalarUDF` referenced inside a `LogicalPlan` +// or an `ExecutionPlan` round-trips identically. +// ============================================================================= + +/// Encode a Python scalar UDF inline if `node` is one. Returns +/// `Ok(true)` when the payload (`DFPYUDF` family prefix, version byte, +/// cloudpickled tuple) was written and the caller should skip its +/// inner codec. Returns `Ok(false)` for any non-Python UDF, signalling +/// the caller to delegate to its `inner`. +pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec) -> Result { + let Some(py_udf) = node.inner().downcast_ref::() else { + return Ok(false); + }; + + Python::attach(|py| -> Result { + let py_version = current_python_version(py) + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; + let bytes = encode_python_scalar_udf(py, py_udf) + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; + write_wire_header(buf, PY_SCALAR_UDF_FAMILY, py_version); + buf.extend_from_slice(&bytes); + Ok(true) + }) +} + +/// Decode an inline Python scalar UDF payload. Returns `Ok(None)` +/// when `buf` does not carry the `DFPYUDF` family prefix, signalling +/// the caller to delegate to its `inner` codec (and eventually the +/// `FunctionRegistry`). +pub(crate) fn try_decode_python_scalar_udf(buf: &[u8]) -> Result>> { + Python::attach(|py| -> Result>> { + let py_version = current_python_version(py) + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; + let Some(payload) = strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", py_version)? + else { + return Ok(None); + }; + let udf = decode_python_scalar_udf(py, payload) + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; + Ok(Some(Arc::new(ScalarUDF::new_from_impl(udf)))) + }) +} + +/// Build the cloudpickle payload for a `PythonFunctionScalarUDF`. +/// +/// Layout: `cloudpickle.dumps((name, func, input_schema_bytes, +/// return_schema_bytes, volatility_str))`. Schema blobs are produced +/// by arrow-rs's native IPC stream writer (no pyarrow round-trip) and +/// decoded with the matching stream reader on the receiver. See +/// [`build_input_schema_bytes`] for what the input blob carries. +fn encode_python_scalar_udf(py: Python<'_>, udf: &PythonFunctionScalarUDF) -> PyResult> { + let signature = udf.signature(); + let input_dtypes = signature_input_dtypes(signature, "PythonFunctionScalarUDF")?; + let input_schema_bytes = build_input_schema_bytes(&input_dtypes)?; + let return_schema_bytes = build_single_field_schema_bytes(udf.return_field().as_ref())?; + let volatility = volatility_wire_str(signature.volatility); + + let payload = PyTuple::new( + py, + [ + udf.name().into_pyobject(py)?.into_any(), + udf.func().bind(py).clone().into_any(), + PyBytes::new(py, &input_schema_bytes).into_any(), + PyBytes::new(py, &return_schema_bytes).into_any(), + volatility.into_pyobject(py)?.into_any(), + ], + )?; + + cloudpickle(py)? + .call_method1("dumps", (payload,))? + .extract::>() +} + +/// Inverse of [`encode_python_scalar_udf`]. +fn decode_python_scalar_udf(py: Python<'_>, payload: &[u8]) -> PyResult { + let tuple = cloudpickle(py)? + .call_method1("loads", (PyBytes::new(py, payload),))? + .cast_into::()?; + + let name: String = tuple.get_item(0)?.extract()?; + let func: Py = tuple.get_item(1)?.unbind(); + let input_schema_bytes: Vec = tuple.get_item(2)?.extract()?; + let return_schema_bytes: Vec = tuple.get_item(3)?.extract()?; + let volatility_str: String = tuple.get_item(4)?.extract()?; + + let input_types = read_input_dtypes(&input_schema_bytes)?; + let return_field = read_single_return_field(&return_schema_bytes, "PythonFunctionScalarUDF")?; + let volatility = parse_volatility_str(&volatility_str)?; + + Ok(PythonFunctionScalarUDF::from_parts( + name, + func, + input_types, + return_field, + volatility, + )) +} + +/// Serialize a `Schema` to a self-contained IPC stream containing +/// only the schema message (no record batches). Inverse: +/// [`schema_from_ipc_bytes`]. +fn schema_to_ipc_bytes(schema: &Schema) -> arrow::error::Result> { + let mut buf: Vec = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, schema)?; + writer.finish()?; + } + Ok(buf) +} + +/// Decode an IPC stream containing only a schema message back into a +/// `Schema`. Inverse: [`schema_to_ipc_bytes`]. +fn schema_from_ipc_bytes(bytes: &[u8]) -> arrow::error::Result { + let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)?; + Ok(reader.schema().as_ref().clone()) +} + +/// Extract the per-arg `DataType`s from a `Signature` known to be +/// `TypeSignature::Exact` (all Python-defined UDFs are constructed +/// with `Signature::exact`). Any other variant indicates the impl was +/// not built by this crate's UDF/UDAF/UDWF constructors. +fn signature_input_dtypes(signature: &Signature, kind: &str) -> PyResult> { + match &signature.type_signature { + TypeSignature::Exact(types) => Ok(types.clone()), + other => Err(pyo3::exceptions::PyValueError::new_err(format!( + "{kind} expected Signature::Exact, got {other:?}" + ))), + } +} + +/// Wrap per-arg `DataType`s in synthetic `arg_{i}` fields and emit +/// the IPC schema blob the encoder writes into the cloudpickle tuple. +/// +/// The names and `nullable: true` are arbitrary: the underlying +/// `TypeSignature::Exact` carries no per-input nullability or +/// metadata, and the receiver collapses these fields back to +/// `Vec` via [`read_input_dtypes`], so anything set here +/// beyond the data type is discarded on decode. +fn build_input_schema_bytes(dtypes: &[DataType]) -> PyResult> { + let fields: Vec = dtypes + .iter() + .enumerate() + .map(|(i, dt)| Field::new(format!("arg_{i}"), dt.clone(), true)) + .collect(); + schema_to_ipc_bytes(&Schema::new(fields)).map_err(arrow_to_py_err) +} + +/// Emit a single-field IPC schema blob. Used for return-type and +/// state-field payloads where the receiver needs to recover field +/// metadata (names, nullability, key/value attributes) verbatim. +fn build_single_field_schema_bytes(field: &Field) -> PyResult> { + schema_to_ipc_bytes(&Schema::new(vec![field.clone()])).map_err(arrow_to_py_err) +} + +/// Decode the per-arg `DataType`s the encoder wrote via +/// [`build_input_schema_bytes`]. +fn read_input_dtypes(bytes: &[u8]) -> PyResult> { + let schema = schema_from_ipc_bytes(bytes).map_err(arrow_to_py_err)?; + Ok(schema + .fields() + .iter() + .map(|f| f.data_type().clone()) + .collect()) +} + +/// Decode a single-field IPC schema blob and return that field by +/// value. `kind` names the UDF flavor in the error message produced +/// when the blob is empty (should be unreachable for sender-side +/// payloads built via [`build_single_field_schema_bytes`]). +fn read_single_return_field(bytes: &[u8], kind: &str) -> PyResult { + let schema = schema_from_ipc_bytes(bytes).map_err(arrow_to_py_err)?; + let field = schema.fields().first().ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err(format!( + "{kind} return schema must contain exactly one field" + )) + })?; + Ok(field.as_ref().clone()) +} + +fn arrow_to_py_err(e: arrow::error::ArrowError) -> PyErr { + pyo3::exceptions::PyValueError::new_err(format!("{e}")) +} + +fn parse_volatility_str(s: &str) -> PyResult { + datafusion_python_util::parse_volatility(s) + .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}"))) +} + +/// Stable wire-format string for a `Volatility`. Pinned to the three +/// tokens [`datafusion_python_util::parse_volatility`] accepts, so an +/// upstream change to `Volatility`'s `Debug` repr cannot silently +/// produce bytes the decoder rejects. +fn volatility_wire_str(v: Volatility) -> &'static str { + match v { + Volatility::Immutable => "immutable", + Volatility::Stable => "stable", + Volatility::Volatile => "volatile", + } +} + +/// Read the interpreter's `sys.version_info` as `(major, minor)`. +/// +/// Used by encoder/decoder to stamp and verify the Python version a +/// cloudpickle payload was produced on. cloudpickle is not portable +/// across Python minor versions; the wire header carries these bytes +/// so a mismatch surfaces an actionable error instead of an opaque +/// `marshal` failure at `cloudpickle.loads` time. +fn current_python_version(py: Python<'_>) -> PyResult<(u8, u8)> { + let version_info = py.import("sys")?.getattr("version_info")?; + let major: u8 = version_info.getattr("major")?.extract()?; + let minor: u8 = version_info.getattr("minor")?.extract()?; + Ok((major, minor)) +} + +/// Cached handle to the `cloudpickle` module. +/// +/// The encode/decode helpers above would otherwise re-resolve the +/// module on every call. `py.import` is backed by `sys.modules` and +/// therefore cheap, but each call still walks a dict and re-binds the +/// result; a plan with many Python UDFs pays that cost per UDF. +/// +/// `PyOnceLock` scopes the cached `Py` to the current +/// interpreter, so the slot drops cleanly on interpreter teardown +/// (relevant under CPython subinterpreters, PEP 684) instead of +/// resurrecting a `Py` rooted in a dead interpreter on the next call. +fn cloudpickle<'py>(py: Python<'py>) -> PyResult> { + static CLOUDPICKLE: PyOnceLock> = PyOnceLock::new(); + CLOUDPICKLE + .get_or_try_init(py, || Ok(py.import("cloudpickle")?.unbind().into_any())) + .map(|cached| cached.bind(py).clone()) +} + +#[cfg(test)] +mod wire_header_tests { + use super::*; + + const TEST_PY: (u8, u8) = (3, 12); + + #[test] + fn strip_returns_none_when_family_absent() { + let buf = b"OTHER_PAYLOAD"; + assert!(matches!( + strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY), + Ok(None) + )); + } + + #[test] + fn strip_errors_on_truncated_version_byte() { + let buf = PY_SCALAR_UDF_FAMILY; + let err = strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err(); + assert!(format!("{err}").contains("missing wire-format version byte")); + } + + #[test] + fn strip_errors_on_too_new_version() { + let mut buf = PY_SCALAR_UDF_FAMILY.to_vec(); + buf.push(WIRE_VERSION_CURRENT.saturating_add(1)); + buf.push(TEST_PY.0); + buf.push(TEST_PY.1); + buf.extend_from_slice(b"payload"); + let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err(); + let msg = format!("{err}"); + assert!(msg.contains("wire-format version v")); + assert!(msg.contains("supports")); + assert!(msg.contains("Align datafusion-python versions")); + } + + #[test] + fn strip_errors_on_too_old_version() { + if WIRE_VERSION_MIN_SUPPORTED == 0 { + return; + } + let mut buf = PY_SCALAR_UDF_FAMILY.to_vec(); + buf.push(WIRE_VERSION_MIN_SUPPORTED - 1); + buf.push(TEST_PY.0); + buf.push(TEST_PY.1); + buf.extend_from_slice(b"payload"); + assert!(strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).is_err()); + } + + #[test] + fn strip_errors_on_truncated_py_major() { + let mut buf = PY_SCALAR_UDF_FAMILY.to_vec(); + buf.push(WIRE_VERSION_CURRENT); + let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err(); + assert!(format!("{err}").contains("missing Python major version byte")); + } + + #[test] + fn strip_errors_on_truncated_py_minor() { + let mut buf = PY_SCALAR_UDF_FAMILY.to_vec(); + buf.push(WIRE_VERSION_CURRENT); + buf.push(TEST_PY.0); + let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err(); + assert!(format!("{err}").contains("missing Python minor version byte")); + } + + #[test] + fn strip_errors_on_py_minor_mismatch() { + let mut buf = Vec::new(); + write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, (3, 11)); + buf.extend_from_slice(b"payload"); + let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", (3, 12)).unwrap_err(); + let msg = format!("{err}"); + assert!(msg.contains("Python 3.11")); + assert!(msg.contains("Python 3.12")); + assert!(msg.contains("not portable across Python minor versions")); + } + + #[test] + fn strip_errors_on_py_major_mismatch() { + let mut buf = Vec::new(); + write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, (3, 12)); + buf.extend_from_slice(b"payload"); + assert!(strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", (4, 0)).is_err()); + } + + #[test] + fn write_then_strip_round_trips_payload() { + let mut buf = Vec::new(); + write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, TEST_PY); + buf.extend_from_slice(b"scalar-payload"); + + let payload = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY) + .unwrap() + .unwrap(); + assert_eq!(payload, b"scalar-payload"); + } +} diff --git a/crates/core/src/udf.rs b/crates/core/src/udf.rs index d48bc729c..2006401db 100644 --- a/crates/core/src/udf.rs +++ b/crates/core/src/udf.rs @@ -42,7 +42,7 @@ use crate::expr::PyExpr; /// This struct holds the Python written function that is a /// ScalarUDF. #[derive(Debug)] -struct PythonFunctionScalarUDF { +pub(crate) struct PythonFunctionScalarUDF { name: String, func: Py, signature: Signature, @@ -66,6 +66,37 @@ impl PythonFunctionScalarUDF { return_field: Arc::new(return_field), } } + + /// Stored Python callable. Consumed by the codec to cloudpickle + /// the function body across process boundaries. + pub(crate) fn func(&self) -> &Py { + &self.func + } + + pub(crate) fn return_field(&self) -> &FieldRef { + &self.return_field + } + + /// Reconstruct a `PythonFunctionScalarUDF` from the parts emitted + /// by the codec. Inputs collapse to `Vec` because + /// `Signature::exact` cannot carry per-input nullability or + /// metadata — the encoder is free to discard that side of the + /// schema. `return_field` is kept as a `Field` so the post-decode + /// nullability and metadata match the sender's instance. + pub(crate) fn from_parts( + name: String, + func: Py, + input_types: Vec, + return_field: Field, + volatility: Volatility, + ) -> Self { + Self { + name, + func, + signature: Signature::exact(input_types, volatility), + return_field: Arc::new(return_field), + } + } } impl Eq for PythonFunctionScalarUDF {} @@ -74,21 +105,51 @@ impl PartialEq for PythonFunctionScalarUDF { self.name == other.name && self.signature == other.signature && self.return_field == other.return_field - && Python::attach(|py| self.func.bind(py).eq(other.func.bind(py)).unwrap_or(false)) + // Identical pointers ⇒ same Python object. Most equality + // checks compare `Arc`-shared clones of the same UDF + // (e.g. expression rewriting), so the pointer match short- + // circuits before touching the GIL. + && (self.func.as_ptr() == other.func.as_ptr() + || Python::attach(|py| { + // Rust's `PartialEq` cannot return `Result`, so we + // have to pick a side when Python `__eq__` raises. + // `false` is the conservative choice — better to + // report two UDFs as distinct than to wrongly + // merge them — but the silent miss can still + // surface as expression-dedup or cache-lookup + // anomalies. Log at `debug` so the failure is + // observable without flooding production logs. + // FIXME: revisit if upstream `ScalarUDFImpl` + // exposes a fallible `PartialEq`. + self.func + .bind(py) + .eq(other.func.bind(py)) + .unwrap_or_else(|e| { + log::debug!( + target: "datafusion_python::udf", + "PythonFunctionScalarUDF {:?} __eq__ raised; treating as unequal: {e}", + self.name, + ); + false + }) + })) } } impl Hash for PythonFunctionScalarUDF { fn hash(&self, state: &mut H) { + // Hash only the identifying header (name + signature + return + // field). Skipping `func` is intentional: the Rust `Hash` + // contract requires `a == b ⇒ hash(a) == hash(b)`, not the + // converse, so a coarser hash is sound — `PartialEq` still + // disambiguates two UDFs with the same header but distinct + // callables. Falling back to a sentinel on `py_hash` failure + // (as a prior revision did) silently mapped every unhashable + // closure to the same bucket; that is the worst case for a + // hashmap and is what this rewrite avoids. self.name.hash(state); self.signature.hash(state); self.return_field.hash(state); - - Python::attach(|py| { - let py_hash = self.func.bind(py).hash().unwrap_or(0); // Handle unhashable objects - - state.write_isize(py_hash); - }); } } @@ -215,4 +276,9 @@ impl PyScalarUDF { fn __repr__(&self) -> PyResult { Ok(format!("ScalarUDF({})", self.function.name())) } + + #[getter] + fn name(&self) -> &str { + self.function.name() + } } diff --git a/pyproject.toml b/pyproject.toml index 951f7adc3..a02f4608a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,13 @@ classifiers = [ "Programming Language :: Rust", ] dependencies = [ + # cloudpickle is invoked by the Rust-side PythonLogicalCodec / + # PythonPhysicalCodec via pyo3 to serialize Python UDF callables — + # scalar, aggregate, and window — into the proto wire format. + # Lazy-imported on the encode / decode hot paths (and cached after + # the first import), so users who never serialize a plan or + # expression incur no runtime cost beyond the install footprint. + "cloudpickle>=2.0", "pyarrow>=16.0.0;python_version<'3.14'", "pyarrow>=22.0.0;python_version>='3.14'", "typing-extensions;python_version<'3.13'", diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 601419fab..f8d523e0d 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -65,7 +65,7 @@ import importlib_metadata # type: ignore[import] # Public submodules -from . import functions, object_store, substrait, unparser +from . import functions, ipc, object_store, substrait, unparser # The following imports are okay to remain as opaque to the user. from ._internal import Config @@ -149,6 +149,7 @@ "configure_formatter", "expr", "functions", + "ipc", "lit", "literal", "object_store", diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 55cc2e52a..645bd9c18 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -46,7 +46,7 @@ from __future__ import annotations -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from typing import TYPE_CHECKING, Any, ClassVar import pyarrow as pa @@ -440,23 +440,165 @@ def variant_name(self) -> str: return self.expr.variant_name() def to_bytes(self, ctx: SessionContext | None = None) -> bytes: - """Serialize this expression to protobuf bytes. + """Serialize this expression to bytes for shipping to another process. - When ``ctx`` is supplied, encoding routes through the session's - installed :class:`LogicalExtensionCodec`. Without ``ctx`` a - default codec is used. + Use this — or :func:`pickle.dumps` — to send an expression to a + worker process for distributed evaluation. + + When ``ctx`` is supplied, encoding routes through that session's + installed :class:`LogicalExtensionCodec`. When ``ctx`` is + ``None``, the default codec is used. + + Built-in functions and Python scalar UDFs travel inside the + returned bytes; the worker does not need to pre-register them. + UDFs imported via the FFI capsule protocol travel by name only + and must be registered on the worker. + + .. warning:: Security + Bytes returned here may embed a cloudpickled Python + callable (when the expression carries a Python scalar UDF). + Reconstructing them via :meth:`from_bytes` or + :func:`pickle.loads` executes arbitrary Python on the + receiver. Only accept payloads from trusted sources. + + .. warning:: Portability + cloudpickle serializes Python bytecode, which is **not + stable across Python minor versions**. A payload produced + on Python 3.11 will fail to load on Python 3.12. The + wire format stamps the sender's ``(major, minor)``; + :meth:`from_bytes` raises a :class:`ValueError` naming + both versions on mismatch. + + cloudpickle captures the UDF callable **by value** — + bytecode and closure cells inlined — but names the + callable resolves via ``import`` are captured **by + reference** (module path only) and must be importable on + the receiver. + + **Self-contained — works anywhere:** + + .. code-block:: python + + # Lambda: bytecode captured inline + udf(lambda x: x * 2, [pa.int64()], pa.int64(), + volatility="immutable") + + # Locally-defined function: bytecode captured inline + def double(x): + return x * 2 + udf(double, [pa.int64()], pa.int64(), volatility="immutable") + + # Closure over a local variable: value captured inline + factor = 3 + udf(lambda x: x * factor, [pa.int64()], pa.int64(), + volatility="immutable") + + **Requires matching environment on receiver:** + + .. code-block:: python + + # Top-level import: `foo` must be installed on receiver + from foo import double + udf(double, [pa.int64()], pa.int64(), volatility="immutable") + + # Bound method of an imported class: same caveat + from mylib import Transformer + t = Transformer() + udf(t.transform, [pa.int64()], pa.int64(), + volatility="immutable") + + Examples: + >>> from datafusion import col, lit + >>> blob = (col("a") + lit(1)).to_bytes() + >>> isinstance(blob, bytes) + True """ ctx_arg = ctx.ctx if ctx is not None else None return self.expr.to_bytes(ctx_arg) - @staticmethod - def from_bytes(ctx: SessionContext, data: bytes) -> Expr: - """Decode an expression from serialized protobuf bytes. + @classmethod + def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr: + """Reconstruct an expression from serialized bytes. + + Accepts output of :meth:`to_bytes` or :func:`pickle.dumps`. + ``ctx`` is the :class:`SessionContext` used to resolve any + function references that travel by name (e.g. FFI UDFs). When + ``ctx`` is ``None`` the worker context installed via + :func:`datafusion.ipc.set_worker_ctx` is consulted; if no worker + context is installed, the global :class:`SessionContext` is used + (sufficient for built-ins and Python scalar UDFs, plus any UDFs + registered on the global context). + + .. warning:: Security + Decoding may invoke ``cloudpickle.loads`` on bytes embedded + in the payload, which executes arbitrary Python code. Treat + ``buf`` as code, not data — only decode bytes you produced + yourself or received from a trusted sender. + + .. warning:: Portability + cloudpickle payloads are **not portable across Python + minor versions**. The wire format stamps the sender's + ``(major, minor)``; if it does not match the current + interpreter, this method raises :class:`ValueError` + naming both versions. Modules the UDF imports must also + be importable on the receiver — see :meth:`to_bytes` for + by-value vs. by-reference details. + + Examples: + >>> from datafusion import Expr, col, lit + >>> blob = (col("a") + lit(1)).to_bytes() + >>> Expr.from_bytes(blob).canonical_name() + 'a + Int64(1)' + """ + from datafusion.ipc import _resolve_ctx + + resolved = _resolve_ctx(ctx) + return cls(expr_internal.RawExpr.from_bytes(resolved.ctx, buf)) + + def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]: + """Pickle protocol hook. + + Lets expressions be shipped to worker processes via + :func:`pickle.dumps` / :func:`pickle.loads`. Built-in functions + and Python scalar UDFs travel inside the pickle bytes; only + FFI-capsule UDFs require pre-registration on the worker. The + worker's :class:`SessionContext` for resolving those references + is looked up via :func:`datafusion.ipc.set_worker_ctx`, falling + back to the global :class:`SessionContext` if none has been + installed on the worker. + + .. warning:: Security + :func:`pickle.loads` on the returned tuple executes + arbitrary Python on the receiver, including any + cloudpickled UDF callable embedded in the payload. Only + unpickle expressions from trusted sources. + + .. warning:: Portability + Sender and receiver must run the same Python + ``(major, minor)`` version; cloudpickle bytecode is not + portable across minor versions. See :meth:`to_bytes` for + details on what travels by value vs. by reference. + + Examples: + >>> import pickle + >>> from datafusion import col, lit + >>> e = col("a") * lit(2) + >>> pickle.loads(pickle.dumps(e)).canonical_name() + 'a * Int64(2)' + """ + return (Expr._reconstruct, (self.to_bytes(),)) + + @classmethod + def _reconstruct(cls, proto_bytes: bytes) -> Expr: + """Internal entry point used by :meth:`__reduce__` on unpickle. - ``ctx`` provides the function registry for resolving UDF - references and the logical codec for in-band Python payloads. + Examples: + >>> from datafusion import Expr, col, lit + >>> blob = (col("a") + lit(1)).to_bytes() + >>> Expr._reconstruct(blob).canonical_name() + 'a + Int64(1)' """ - return Expr(expr_internal.RawExpr.from_bytes(ctx.ctx, data)) + return cls.from_bytes(proto_bytes) def __richcmp__(self, other: Expr, op: int) -> Expr: """Comparison operator.""" diff --git a/python/datafusion/ipc.py b/python/datafusion/ipc.py new file mode 100644 index 000000000..78b6873f7 --- /dev/null +++ b/python/datafusion/ipc.py @@ -0,0 +1,161 @@ +# 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. + +"""Worker-side setup for distributing DataFusion expressions. + +When a :class:`Expr` is shipped to a worker process (e.g. through +:func:`multiprocessing.Pool` or a Ray actor), the worker reconstructs the +expression against a :class:`SessionContext`. If the expression references +UDFs imported via the FFI capsule protocol — or any UDF the worker would +otherwise resolve from its registered functions rather than from inside +the shipped expression — install a configured :class:`SessionContext` +once per worker: + +.. code-block:: python + + from datafusion import SessionContext + from datafusion.ipc import set_worker_ctx + + def init_worker(): + ctx = SessionContext() + ctx.register_udaf(my_ffi_aggregate) + set_worker_ctx(ctx) + +Built-in functions and Python scalar UDFs travel inside the shipped +expression itself and do not need pre-registration on the worker. + +.. note:: Serialization model + + Expressions containing Python scalar UDFs are serialized using + :mod:`cloudpickle`. The callable itself travels **by value** + (bytecode and closure cells inlined), but any names the callable + resolves via ``import`` are captured **by reference** and must be + importable on the receiving worker. + + The serialized payload is stamped with the sender's Python + ``(major, minor)`` version. Loading on a different minor version + raises :class:`ValueError` with an actionable message — cloudpickle + payloads are not portable across Python minor versions. See + :meth:`datafusion.Expr.to_bytes` for examples of what travels by + value vs. by reference. +""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from datafusion.context import SessionContext + + +__all__ = [ + "clear_worker_ctx", + "get_worker_ctx", + "set_worker_ctx", +] + + +_local = threading.local() + + +def set_worker_ctx(ctx: SessionContext) -> None: + """Install this worker's :class:`SessionContext` for shipped expressions. + + Call once per worker — typically from a ``multiprocessing.Pool`` + initializer or a Ray actor ``__init__``. Idempotent: overwrites any + previous value. Stored in a thread-local slot, so each thread within a + worker may install its own context independently. + + Examples: + >>> from datafusion import SessionContext + >>> from datafusion.ipc import set_worker_ctx, get_worker_ctx, clear_worker_ctx + >>> set_worker_ctx(SessionContext()) + >>> get_worker_ctx() is not None + True + >>> clear_worker_ctx() + """ + _local.ctx = ctx + + +def clear_worker_ctx() -> None: + """Remove this worker's installed :class:`SessionContext`. + + After clearing, expressions reconstructed in this worker fall back to + the global :class:`SessionContext` — adequate for built-ins and Python + scalar UDFs, but anything imported via the FFI capsule protocol must + be registered on the global context to resolve. + + Examples: + >>> from datafusion import SessionContext + >>> from datafusion.ipc import set_worker_ctx, clear_worker_ctx, get_worker_ctx + >>> set_worker_ctx(SessionContext()) + >>> clear_worker_ctx() + >>> get_worker_ctx() is None + True + """ + if hasattr(_local, "ctx"): + del _local.ctx + + +def get_worker_ctx() -> SessionContext | None: + """Return this worker's installed :class:`SessionContext`, or ``None``. + + Examples: + >>> from datafusion.ipc import get_worker_ctx, clear_worker_ctx + >>> clear_worker_ctx() + >>> get_worker_ctx() is None + True + """ + return getattr(_local, "ctx", None) + + +def _resolve_ctx( + explicit_ctx: SessionContext | None = None, +) -> SessionContext: + """Resolve a context for Expr reconstruction. + + Priority: explicit argument > worker context > global context. + Falling back to the global :class:`SessionContext` (instead of a + freshly constructed one) preserves any registrations the user has + installed on it. + + Examples: + >>> from datafusion import SessionContext + >>> from datafusion.ipc import _resolve_ctx, clear_worker_ctx + >>> clear_worker_ctx() + >>> isinstance(_resolve_ctx(), SessionContext) + True + >>> ctx = SessionContext() + >>> _resolve_ctx(ctx) is ctx + True + """ + if explicit_ctx is not None: + return explicit_ctx + worker = get_worker_ctx() + if worker is not None: + return worker + # Lazy import: `datafusion/__init__.py` imports `datafusion.ipc` + # before `datafusion.context`, so a module-top import would force + # `datafusion.context` to load mid-init of `datafusion.ipc`. The + # cycle is benign today (context.py only pulls expr.py at module + # scope, neither pulls ipc.py back), but a single new import in + # context.py's transitive deps could turn it into a real cycle. + # Deferring keeps `datafusion.ipc` import-order-independent. + from datafusion.context import SessionContext # noqa: PLC0415 + + return SessionContext.global_ctx() diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index 848ab4cee..d79cf22e8 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -141,6 +141,29 @@ def __init__( name, func, input_fields, return_field, str(volatility) ) + @property + def name(self) -> str: + """Return the registered name of this UDF. + + For UDFs imported via the FFI capsule protocol, this is the + name the capsule itself reports — not the ``name`` argument + passed to the constructor (which is ignored on the FFI path). + + Examples: + >>> import pyarrow as pa + >>> from datafusion import udf + >>> double = udf( + ... lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]), + ... [pa.int64()], + ... pa.int64(), + ... volatility="immutable", + ... name="double", + ... ) + >>> double.name + 'double' + """ + return self._udf.name + def __repr__(self) -> str: """Print a string representation of the Scalar UDF.""" return self._udf.__repr__() diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index cf1fe43e8..485d69624 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -1189,7 +1189,7 @@ def test_expr_to_bytes_roundtrip(ctx: SessionContext) -> None: original = col("a") + lit(1) blob = original.to_bytes(ctx) - restored = Expr.from_bytes(ctx, blob) + restored = Expr.from_bytes(blob, ctx=ctx) # Canonical name preserves the structure of the expression even # though the underlying PyExpr instances are different. @@ -1204,6 +1204,6 @@ def test_expr_to_bytes_no_ctx_default_codec() -> None: fresh = SessionContext() original = col("a") * lit(2) blob = original.to_bytes() # encode side: default codec - restored = Expr.from_bytes(fresh, blob) + restored = Expr.from_bytes(blob, ctx=fresh) assert restored.canonical_name() == original.canonical_name() diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py new file mode 100644 index 000000000..5d8d9285f --- /dev/null +++ b/python/tests/test_pickle_expr.py @@ -0,0 +1,184 @@ +# 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. + +"""In-process pickle round-trip tests for :class:`Expr`. + +Built-in functions and Python scalar UDFs travel with the pickled +expression and do not need worker-side pre-registration. The worker +context (:mod:`datafusion.ipc`) is only consulted for UDFs imported +via the FFI capsule protocol. +""" + +from __future__ import annotations + +import pickle + +import pyarrow as pa +import pytest +from datafusion import Expr, SessionContext, col, lit, udf +from datafusion.ipc import ( + clear_worker_ctx, + set_worker_ctx, +) + + +@pytest.fixture(autouse=True) +def _reset_worker_ctx(): + """Ensure every test starts with no worker context installed.""" + clear_worker_ctx() + yield + clear_worker_ctx() + + +def _double_udf(): + return udf( + lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]), + [pa.int64()], + pa.int64(), + volatility="immutable", + name="double", + ) + + +class TestProtoRoundTrip: + def test_builtin_round_trip(self): + e = col("a") + lit(1) + blob = pickle.dumps(e) + decoded = pickle.loads(blob) # noqa: S301 + assert decoded.canonical_name() == e.canonical_name() + + def test_to_bytes_from_bytes(self): + e = col("x") * lit(7) + blob = e.to_bytes() + assert isinstance(blob, bytes) + decoded = Expr.from_bytes(blob) + assert decoded.canonical_name() == e.canonical_name() + + def test_explicit_ctx_used(self, ctx): + e = col("a") + lit(1) + decoded = Expr.from_bytes(e.to_bytes(), ctx=ctx) + assert decoded.canonical_name() == e.canonical_name() + + +class TestUDFCodec: + """Python scalar UDFs ride inside the proto blob via the Rust codec. + + No worker context needed on the receiver — the cloudpickled callable is + embedded in ``fun_definition`` and reconstructed automatically. + """ + + def test_udf_self_contained_blob(self): + e = _double_udf()(col("a")) + blob = pickle.dumps(e) + # The codec inlines the callable, so the blob is much bigger than a + # pure built-in blob but doesn't depend on receiver-side registration. + assert len(blob) > 200 + + def test_udf_decodes_into_fresh_ctx(self): + e = _double_udf()(col("a")) + blob = e.to_bytes() + fresh = SessionContext() + decoded = Expr.from_bytes(blob, ctx=fresh) + assert "double" in decoded.canonical_name() + + def test_udf_decodes_via_pickle_with_no_worker_ctx(self): + e = _double_udf()(col("a")) + blob = pickle.dumps(e) + decoded = pickle.loads(blob) # noqa: S301 + assert "double" in decoded.canonical_name() + + def test_udf_decodes_via_pickle_with_worker_ctx(self): + set_worker_ctx(SessionContext()) + e = _double_udf()(col("a")) + blob = pickle.dumps(e) + decoded = pickle.loads(blob) # noqa: S301 + assert "double" in decoded.canonical_name() + + def test_closure_capturing_udf_names_match(self): + captured_multiplier = 7 + + def fn(arr): + return pa.array([(v.as_py() or 0) * captured_multiplier for v in arr]) + + u = udf( + fn, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="times_seven", + ) + e = u(col("a")) + blob = pickle.dumps(e) + decoded = pickle.loads(blob) # noqa: S301 + assert decoded.canonical_name() == e.canonical_name() + + def test_multi_arg_udf_round_trip(self): + """Wire format builds synthetic `arg_{i}` fields per input — exercise + with a 2-arg UDF spanning two distinct DataTypes.""" + add_scaled = udf( + lambda a, b: pa.array( + [ + (x.as_py() or 0) + (y.as_py() or 0.0) + for x, y in zip(a, b, strict=False) + ] + ), + [pa.int64(), pa.float64()], + pa.float64(), + volatility="immutable", + name="add_scaled", + ) + e = add_scaled(col("a"), col("b")) + decoded = pickle.loads(pickle.dumps(e)) # noqa: S301 + assert decoded.canonical_name() == e.canonical_name() + assert "add_scaled" in decoded.canonical_name() + + +class TestErrorPaths: + def test_from_bytes_rejects_garbage(self): + with pytest.raises(Exception): # noqa: B017 + Expr.from_bytes(b"not a valid protobuf payload") + + def test_from_bytes_rejects_empty(self): + with pytest.raises(Exception): # noqa: B017 + Expr.from_bytes(b"") + + def test_cross_version_error_message(self): + """Decoding a payload stamped with a different Python minor + version raises a clear, actionable error rather than an opaque + marshal/unpickle failure. + + The wire frame inside the protobuf is: + ``DFPYUDF (7) | version (1) | py_major (1) | py_minor (1) | cloudpickle``. + We locate the frame inside the outer protobuf and patch the + minor byte at offset 9. + """ + import sys + + e = _double_udf()(col("a")) + blob = e.to_bytes() + + idx = blob.find(b"DFPYUDF") + assert idx >= 0, "DFPYUDF frame not found in payload" + + different_minor = (sys.version_info.minor + 1) % 256 + tampered = bytearray(blob) + tampered[idx + 9] = different_minor + + with pytest.raises( + Exception, match="not portable across Python minor versions" + ): + Expr.from_bytes(bytes(tampered)) diff --git a/uv.lock b/uv.lock index 3b7135e32..3fd3eec4b 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14'", @@ -257,6 +257,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/f6/65ecc6878a89bb1c23a086ea335ad4bf21a588990c3f535a227b9eea9108/charset_normalizer-3.4.1-py3-none-any.whl", hash = "sha256:d98b1668f06378c6dbefec3b92299716b931cd4e6061f3c875a71ced1780ab85", size = 49767, upload-time = "2024-12-24T18:12:32.852Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "codespell" version = "2.4.1" @@ -316,6 +325,7 @@ wheels = [ name = "datafusion" source = { editable = "." } dependencies = [ + { name = "cloudpickle" }, { name = "pyarrow" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] @@ -351,6 +361,7 @@ docs = [ [package.metadata] requires-dist = [ + { name = "cloudpickle", specifier = ">=2.0" }, { name = "pyarrow", marker = "python_full_version < '3.14'", specifier = ">=16.0.0" }, { name = "pyarrow", marker = "python_full_version >= '3.14'", specifier = ">=22.0.0" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, From dac9ec6230dba8717d8a0d27de19a141600486b1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 20 May 2026 18:58:32 -0400 Subject: [PATCH 37/83] feat: enable pickling for Python aggregate and window UDFs (#1545) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: inline encoding for Python aggregate and window UDFs Extends the PythonLogicalCodec / PythonPhysicalCodec inline encoding introduced for scalar UDFs to also cover Python-defined aggregate and window UDFs. The cloudpickle tuple shape per family is: DFPYUDA (agg) (name, accumulator_factory, input_schema_bytes, return_schema_bytes, state_schema_bytes, volatility_str) DFPYUDW (window) (name, evaluator_factory, input_schema_bytes, return_schema_bytes, volatility_str) Same wire-framing as scalar (family magic + version byte + cloudpickle blob), same schema serde (arrow-rs native IPC), same cached cloudpickle handle. The agg state schema is encoded as a full IPC schema so the post-decode UDF reports the same names + nullability + metadata as the sender — relevant for accumulators whose StateFieldsArgs consumers key off names rather than positional DataType. Required restructuring two existing UDF impls so the codec can grab the Python callable directly: * udaf.rs: replaces create_udaf + AccumulatorFactoryFunction closure with a named PythonFunctionAggregateUDF that stores the Py accumulator factory. Synthesizes state_{i} field names when the Python constructor passes only Vec; from_parts preserves the full state schema on the decode side. * udwf.rs: renames MultiColumnWindowUDF -> PythonFunctionWindowUDF, drops the PartitionEvaluatorFactory PtrEq wrapper, stores the Py evaluator directly. PartialEq and Hash get the same pointer-identity fast path + debug-log exception handling already on PythonFunctionScalarUDF. User-facing surface: * AggregateUDF.name and WindowUDF.name properties (parallel to the ScalarUDF.name shipped in PR1). * Existing UDAF/UDWF construction paths are unchanged. The per-session with_python_udf_inlining toggle, sender-side context, strict refusal, and user-guide docs land in PRs 3-4 of this series. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: restore pub UDAF/UDWF helpers and document inline encoding Re-export `to_rust_accumulator`, `to_rust_partition_evaluator`, and `PythonFunctionWindowUDF` (with a `MultiColumnWindowUDF` alias) by promoting `udaf` and `udwf` to `pub mod` so prior downstream Rust consumers keep their API surface after the inline-encoding refactor. Adds an end-to-end window UDF pickle round-trip test that runs the decoded evaluator over a real session, mirroring the aggregate test. Documents the cloudpickle-based shipping behavior of Python aggregate and window UDFs in the user-guide aggregations and windows pages. Co-Authored-By: Claude Opus 4.7 (1M context) * fix: address PR #1545 review feedback - Fix CountAcc.merge in pickle test: sum over states[0] (partition counts), not over the list of state fields. The prior implementation only added partition 0's count when merging across partitions. - Drive test_agg_udf_evaluates_after_roundtrip with a two-batch DataFrame so merge actually runs and the round-tripped state-field schema is exercised end-to-end. - Correct PY_AGG_UDF_FAMILY / PY_WINDOW_UDF_FAMILY doc comments and the aggregate block comment to reference "return schema bytes" rather than "return type" / "return_type_bytes" so the docs match the actual on-wire layout. - Keep `udaf` and `udwf` modules private (matching `udf`) and selectively re-export the helpers downstream Rust consumers rely on (`to_rust_accumulator`, `to_rust_partition_evaluator`, `PythonFunctionWindowUDF`, `MultiColumnWindowUDF`) instead of exposing the whole module surface. - Rename codec helpers `*_agg_udf` -> `*_udaf` and `*_window_udf` -> `*_udwf` for naming consistency with the Python public aliases. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/codec.rs | 316 +++++++++++++++++- crates/core/src/lib.rs | 5 + crates/core/src/udaf.rs | 190 ++++++++++- crates/core/src/udwf.rs | 126 +++++-- .../common-operations/aggregations.rst | 18 + .../user-guide/common-operations/windows.rst | 18 + python/datafusion/expr.py | 25 +- python/datafusion/ipc.py | 19 +- python/datafusion/user_defined.py | 20 ++ python/tests/test_pickle_expr.py | 137 +++++++- 10 files changed, 787 insertions(+), 87 deletions(-) diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index cc038edc9..363ee82b8 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -66,12 +66,14 @@ //! | Layer + kind | Family prefix | //! | ----------------------------- | ------------- | //! | `PythonLogicalCodec` scalar | `DFPYUDF` | +//! | `PythonLogicalCodec` agg | `DFPYUDA` | +//! | `PythonLogicalCodec` window | `DFPYUDW` | //! | `PythonPhysicalCodec` scalar | `DFPYUDF` | +//! | `PythonPhysicalCodec` agg | `DFPYUDA` | +//! | `PythonPhysicalCodec` window | `DFPYUDW` | //! | User FFI extension codec | user-chosen | //! | Default codec | (none) | //! -//! Aggregate and window UDF families are reserved for follow-on work. -//! //! Current wire-format version is [`WIRE_VERSION_CURRENT`]; supported //! receive range is `WIRE_VERSION_MIN_SUPPORTED..=WIRE_VERSION_CURRENT`. //! Bump [`WIRE_VERSION_CURRENT`] whenever the cloudpickle tuple shape @@ -94,8 +96,8 @@ use datafusion::datasource::TableProvider; use datafusion::datasource::file_format::FileFormatFactory; use datafusion::execution::TaskContext; use datafusion::logical_expr::{ - AggregateUDF, Extension, LogicalPlan, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, - Volatility, WindowUDF, + AggregateUDF, AggregateUDFImpl, Extension, LogicalPlan, ScalarUDF, ScalarUDFImpl, Signature, + TypeSignature, Volatility, WindowUDF, WindowUDFImpl, }; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; @@ -105,7 +107,10 @@ use pyo3::prelude::*; use pyo3::sync::PyOnceLock; use pyo3::types::{PyBytes, PyTuple}; +use crate::errors::to_datafusion_err; +use crate::udaf::PythonFunctionAggregateUDF; use crate::udf::PythonFunctionScalarUDF; +use crate::udwf::PythonFunctionWindowUDF; // Wire-format framing for inlined Python UDF payloads. // @@ -126,6 +131,17 @@ use crate::udf::PythonFunctionScalarUDF; /// volatility). pub(crate) const PY_SCALAR_UDF_FAMILY: &[u8] = b"DFPYUDF"; +/// Family prefix for an inlined Python aggregate UDF +/// (cloudpickled tuple of name, accumulator factory, input schema bytes, +/// return schema bytes (single-field IPC schema), state schema bytes, +/// volatility). +pub(crate) const PY_AGG_UDF_FAMILY: &[u8] = b"DFPYUDA"; + +/// Family prefix for an inlined Python window UDF +/// (cloudpickled tuple of name, evaluator factory, input schema bytes, +/// return schema bytes (single-field IPC schema), volatility). +pub(crate) const PY_WINDOW_UDF_FAMILY: &[u8] = b"DFPYUDW"; + /// Wire-format version this build emits. pub(crate) const WIRE_VERSION_CURRENT: u8 = 1; @@ -299,18 +315,30 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { + if try_encode_python_udaf(node, buf)? { + return Ok(()); + } self.inner.try_encode_udaf(node, buf) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + if let Some(udaf) = try_decode_python_udaf(buf)? { + return Ok(udaf); + } self.inner.try_decode_udaf(name, buf) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + if try_encode_python_udwf(node, buf)? { + return Ok(()); + } self.inner.try_encode_udwf(node, buf) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + if let Some(udwf) = try_decode_python_udwf(buf)? { + return Ok(udwf); + } self.inner.try_decode_udwf(name, buf) } } @@ -389,18 +417,30 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { + if try_encode_python_udaf(node, buf)? { + return Ok(()); + } self.inner.try_encode_udaf(node, buf) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + if let Some(udaf) = try_decode_python_udaf(buf)? { + return Ok(udaf); + } self.inner.try_decode_udaf(name, buf) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + if try_encode_python_udwf(node, buf)? { + return Ok(()); + } self.inner.try_encode_udwf(node, buf) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + if let Some(udwf) = try_decode_python_udwf(buf)? { + return Ok(udwf); + } self.inner.try_decode_udwf(name, buf) } } @@ -425,12 +465,8 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec) }; Python::attach(|py| -> Result { - let py_version = current_python_version(py) - .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; - let bytes = encode_python_scalar_udf(py, py_udf) - .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; - write_wire_header(buf, PY_SCALAR_UDF_FAMILY, py_version); - buf.extend_from_slice(&bytes); + let bytes = encode_python_scalar_udf(py, py_udf).map_err(to_datafusion_err)?; + append_framed_payload(py, buf, PY_SCALAR_UDF_FAMILY, &bytes)?; Ok(true) }) } @@ -441,14 +477,11 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec) /// `FunctionRegistry`). pub(crate) fn try_decode_python_scalar_udf(buf: &[u8]) -> Result>> { Python::attach(|py| -> Result>> { - let py_version = current_python_version(py) - .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; - let Some(payload) = strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", py_version)? + let Some(payload) = read_framed_payload(py, buf, PY_SCALAR_UDF_FAMILY, "scalar UDF")? else { return Ok(None); }; - let udf = decode_python_scalar_udf(py, payload) - .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; + let udf = decode_python_scalar_udf(py, payload).map_err(to_datafusion_err)?; Ok(Some(Arc::new(ScalarUDF::new_from_impl(udf)))) }) } @@ -564,6 +597,11 @@ fn build_single_field_schema_bytes(field: &Field) -> PyResult> { schema_to_ipc_bytes(&Schema::new(vec![field.clone()])).map_err(arrow_to_py_err) } +/// Emit a multi-field IPC schema blob. +fn build_schema_bytes(fields: Vec) -> PyResult> { + schema_to_ipc_bytes(&Schema::new(fields)).map_err(arrow_to_py_err) +} + /// Decode the per-arg `DataType`s the encoder wrote via /// [`build_input_schema_bytes`]. fn read_input_dtypes(bytes: &[u8]) -> PyResult> { @@ -624,6 +662,37 @@ fn current_python_version(py: Python<'_>) -> PyResult<(u8, u8)> { Ok((major, minor)) } +/// Stamp `buf` with the framing header for `family` plus the current +/// Python `(major, minor)`, then append `payload`. Bundles the +/// `current_python_version` lookup with the header write so each +/// encoder call site stays one line. +fn append_framed_payload( + py: Python<'_>, + buf: &mut Vec, + family: &[u8], + payload: &[u8], +) -> Result<()> { + let py_version = current_python_version(py).map_err(to_datafusion_err)?; + write_wire_header(buf, family, py_version); + buf.extend_from_slice(payload); + Ok(()) +} + +/// Inspect `buf`'s framing against `family` + the current Python +/// `(major, minor)`. Returns `Ok(None)` when `buf` does not carry +/// `family` (caller should delegate); `Ok(Some(payload))` when the +/// framing matches; `Err(_)` for a recognised family at the wrong +/// wire-format or Python version (see [`strip_wire_header`]). +fn read_framed_payload<'a>( + py: Python<'_>, + buf: &'a [u8], + family: &[u8], + kind: &str, +) -> Result> { + let py_version = current_python_version(py).map_err(to_datafusion_err)?; + strip_wire_header(buf, family, kind, py_version) +} + /// Cached handle to the `cloudpickle` module. /// /// The encode/decode helpers above would otherwise re-resolve the @@ -642,6 +711,186 @@ fn cloudpickle<'py>(py: Python<'py>) -> PyResult> { .map(|cached| cached.bind(py).clone()) } +// ============================================================================= +// Shared Python window UDF encode / decode helpers +// +// Cloudpickle tuple shape: `(name, evaluator_factory, input_schema_bytes, +// return_schema_bytes, volatility_str)`. The evaluator factory is the +// Python callable that produces a new evaluator instance per partition. +// ============================================================================= + +pub(crate) fn try_encode_python_udwf(node: &WindowUDF, buf: &mut Vec) -> Result { + let Some(py_udf) = node.inner().downcast_ref::() else { + return Ok(false); + }; + + Python::attach(|py| -> Result { + let bytes = encode_python_udwf(py, py_udf).map_err(to_datafusion_err)?; + append_framed_payload(py, buf, PY_WINDOW_UDF_FAMILY, &bytes)?; + Ok(true) + }) +} + +pub(crate) fn try_decode_python_udwf(buf: &[u8]) -> Result>> { + Python::attach(|py| -> Result>> { + let Some(payload) = read_framed_payload(py, buf, PY_WINDOW_UDF_FAMILY, "window UDF")? + else { + return Ok(None); + }; + let udf = decode_python_udwf(py, payload).map_err(to_datafusion_err)?; + Ok(Some(Arc::new(WindowUDF::new_from_impl(udf)))) + }) +} + +fn encode_python_udwf(py: Python<'_>, udf: &PythonFunctionWindowUDF) -> PyResult> { + let signature = WindowUDFImpl::signature(udf); + let input_dtypes = signature_input_dtypes(signature, "PythonFunctionWindowUDF")?; + let input_schema_bytes = build_input_schema_bytes(&input_dtypes)?; + let return_field = Field::new("result", udf.return_type().clone(), true); + let return_schema_bytes = build_single_field_schema_bytes(&return_field)?; + let volatility = volatility_wire_str(signature.volatility); + + let payload = PyTuple::new( + py, + [ + WindowUDFImpl::name(udf).into_pyobject(py)?.into_any(), + udf.evaluator().bind(py).clone().into_any(), + PyBytes::new(py, &input_schema_bytes).into_any(), + PyBytes::new(py, &return_schema_bytes).into_any(), + volatility.into_pyobject(py)?.into_any(), + ], + )?; + + cloudpickle(py)? + .call_method1("dumps", (payload,))? + .extract::>() +} + +fn decode_python_udwf(py: Python<'_>, payload: &[u8]) -> PyResult { + let tuple = cloudpickle(py)? + .call_method1("loads", (PyBytes::new(py, payload),))? + .cast_into::()?; + + let name: String = tuple.get_item(0)?.extract()?; + let evaluator: Py = tuple.get_item(1)?.unbind(); + let input_schema_bytes: Vec = tuple.get_item(2)?.extract()?; + let return_schema_bytes: Vec = tuple.get_item(3)?.extract()?; + let volatility_str: String = tuple.get_item(4)?.extract()?; + + let input_types = read_input_dtypes(&input_schema_bytes)?; + let return_type = read_single_return_field(&return_schema_bytes, "PythonFunctionWindowUDF")? + .data_type() + .clone(); + let volatility = parse_volatility_str(&volatility_str)?; + + Ok(PythonFunctionWindowUDF::new( + name, + evaluator, + input_types, + return_type, + volatility, + )) +} + +// ============================================================================= +// Shared Python aggregate UDF encode / decode helpers +// +// Cloudpickle tuple shape: `(name, accumulator_factory, input_schema_bytes, +// return_schema_bytes, state_schema_bytes, volatility_str)`. The accumulator +// factory is the Python callable that produces a new accumulator instance +// per partition. +// ============================================================================= + +pub(crate) fn try_encode_python_udaf(node: &AggregateUDF, buf: &mut Vec) -> Result { + let Some(py_udf) = node.inner().downcast_ref::() else { + return Ok(false); + }; + + Python::attach(|py| -> Result { + let bytes = encode_python_udaf(py, py_udf).map_err(to_datafusion_err)?; + append_framed_payload(py, buf, PY_AGG_UDF_FAMILY, &bytes)?; + Ok(true) + }) +} + +pub(crate) fn try_decode_python_udaf(buf: &[u8]) -> Result>> { + Python::attach(|py| -> Result>> { + let Some(payload) = read_framed_payload(py, buf, PY_AGG_UDF_FAMILY, "aggregate UDF")? + else { + return Ok(None); + }; + let udf = decode_python_udaf(py, payload).map_err(to_datafusion_err)?; + Ok(Some(Arc::new(AggregateUDF::new_from_impl(udf)))) + }) +} + +fn encode_python_udaf(py: Python<'_>, udf: &PythonFunctionAggregateUDF) -> PyResult> { + let signature = AggregateUDFImpl::signature(udf); + let input_dtypes = signature_input_dtypes(signature, "PythonFunctionAggregateUDF")?; + let input_schema_bytes = build_input_schema_bytes(&input_dtypes)?; + let return_field = Field::new("result", udf.return_type().clone(), true); + let return_schema_bytes = build_single_field_schema_bytes(&return_field)?; + let state_fields: Vec = udf + .state_fields_ref() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + let state_schema_bytes = build_schema_bytes(state_fields)?; + let volatility = volatility_wire_str(signature.volatility); + + let payload = PyTuple::new( + py, + [ + AggregateUDFImpl::name(udf).into_pyobject(py)?.into_any(), + udf.accumulator().bind(py).clone().into_any(), + PyBytes::new(py, &input_schema_bytes).into_any(), + PyBytes::new(py, &return_schema_bytes).into_any(), + PyBytes::new(py, &state_schema_bytes).into_any(), + volatility.into_pyobject(py)?.into_any(), + ], + )?; + + cloudpickle(py)? + .call_method1("dumps", (payload,))? + .extract::>() +} + +fn decode_python_udaf(py: Python<'_>, payload: &[u8]) -> PyResult { + let tuple = cloudpickle(py)? + .call_method1("loads", (PyBytes::new(py, payload),))? + .cast_into::()?; + + let name: String = tuple.get_item(0)?.extract()?; + let accumulator: Py = tuple.get_item(1)?.unbind(); + let input_schema_bytes: Vec = tuple.get_item(2)?.extract()?; + let return_schema_bytes: Vec = tuple.get_item(3)?.extract()?; + let state_schema_bytes: Vec = tuple.get_item(4)?.extract()?; + let volatility_str: String = tuple.get_item(5)?.extract()?; + + let input_types = read_input_dtypes(&input_schema_bytes)?; + let return_type = read_single_return_field(&return_schema_bytes, "PythonFunctionAggregateUDF")? + .data_type() + .clone(); + // Preserve the encoded state field metadata (names, nullability, + // arbitrary key/value attributes) so the post-decode UDF reports + // the same state schema as the sender's instance — important for + // accumulators whose `StateFieldsArgs` consumers key off names or + // nullability rather than positional `DataType`. + let state_schema = schema_from_ipc_bytes(&state_schema_bytes).map_err(arrow_to_py_err)?; + let state_fields: Vec = + state_schema.fields().iter().cloned().collect(); + let volatility = parse_volatility_str(&volatility_str)?; + + Ok(PythonFunctionAggregateUDF::from_parts( + name, + accumulator, + input_types, + return_type, + state_fields, + volatility, + )) +} + #[cfg(test)] mod wire_header_tests { use super::*; @@ -729,7 +978,7 @@ mod wire_header_tests { } #[test] - fn write_then_strip_round_trips_payload() { + fn write_then_strip_round_trips_scalar_payload() { let mut buf = Vec::new(); write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, TEST_PY); buf.extend_from_slice(b"scalar-payload"); @@ -739,4 +988,39 @@ mod wire_header_tests { .unwrap(); assert_eq!(payload, b"scalar-payload"); } + + #[test] + fn write_then_strip_round_trips_agg_payload() { + let mut buf = Vec::new(); + write_wire_header(&mut buf, PY_AGG_UDF_FAMILY, TEST_PY); + buf.extend_from_slice(b"agg-payload"); + + let payload = strip_wire_header(&buf, PY_AGG_UDF_FAMILY, "aggregate UDF", TEST_PY) + .unwrap() + .unwrap(); + assert_eq!(payload, b"agg-payload"); + } + + #[test] + fn write_then_strip_round_trips_window_payload() { + let mut buf = Vec::new(); + write_wire_header(&mut buf, PY_WINDOW_UDF_FAMILY, TEST_PY); + buf.extend_from_slice(b"window-payload"); + + let payload = strip_wire_header(&buf, PY_WINDOW_UDF_FAMILY, "window UDF", TEST_PY) + .unwrap() + .unwrap(); + assert_eq!(payload, b"window-payload"); + } + + #[test] + fn strip_does_not_match_a_different_family() { + let mut buf = Vec::new(); + write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, TEST_PY); + buf.extend_from_slice(b"payload"); + assert!(matches!( + strip_wire_header(&buf, PY_WINDOW_UDF_FAMILY, "window UDF", TEST_PY), + Ok(None) + )); + } } diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index e3551c937..8b622d344 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -65,6 +65,11 @@ mod udf; pub mod udtf; mod udwf; +// Re-export helpers previously consumed by downstream Rust crates. +// Modules stay private to keep the public Rust API surface small. +pub use udaf::to_rust_accumulator; +pub use udwf::{MultiColumnWindowUDF, PythonFunctionWindowUDF, to_rust_partition_evaluator}; + #[cfg(feature = "mimalloc")] #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; diff --git a/crates/core/src/udaf.rs b/crates/core/src/udaf.rs index 80ef51716..caf7b97bc 100644 --- a/crates/core/src/udaf.rs +++ b/crates/core/src/udaf.rs @@ -19,12 +19,13 @@ use std::ptr::NonNull; use std::sync::Arc; use datafusion::arrow::array::ArrayRef; -use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow}; use datafusion::common::ScalarValue; use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion::logical_expr::{ - Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, create_udaf, + Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, Signature, Volatility, }; use datafusion_ffi::udaf::FFI_AggregateUDF; use datafusion_python_util::parse_volatility; @@ -144,15 +145,168 @@ impl Accumulator for RustAccumulator { } } +fn instantiate_accumulator(accum: &Py) -> Result> { + let instance = Python::attach(|py| { + accum + .call0(py) + .map_err(|e| DataFusionError::Execution(format!("{e}"))) + })?; + Ok(Box::new(RustAccumulator::new(instance))) +} + +/// Wrap a Python accumulator factory in an `AccumulatorFactoryFunction`. +/// +/// Retained for downstream callers that previously consumed this +/// helper to build a [`AccumulatorFactoryFunction`] for `create_udaf` +/// or similar factory-based APIs. New in-crate code should construct +/// a [`PythonFunctionAggregateUDF`] directly so the codec can downcast +/// and ship it inline. pub fn to_rust_accumulator(accum: Py) -> AccumulatorFactoryFunction { - Arc::new(move |_args| -> Result> { - let accum = Python::attach(|py| { - accum - .call0(py) - .map_err(|e| DataFusionError::Execution(format!("{e}"))) - })?; - Ok(Box::new(RustAccumulator::new(accum))) - }) + Arc::new(move |_args| instantiate_accumulator(&accum)) +} + +/// Named-struct `AggregateUDFImpl` for Python-defined aggregate UDFs. +/// Holds the Python accumulator factory directly so the codec can +/// downcast and cloudpickle it across process boundaries. +#[derive(Debug)] +pub(crate) struct PythonFunctionAggregateUDF { + name: String, + accumulator: Py, + signature: Signature, + return_type: DataType, + state_fields: Vec, +} + +impl PythonFunctionAggregateUDF { + fn new( + name: String, + accumulator: Py, + input_types: Vec, + return_type: DataType, + state_types: Vec, + volatility: Volatility, + ) -> Self { + let signature = Signature::exact(input_types, volatility); + let state_fields = state_types + .into_iter() + .enumerate() + .map(|(i, t)| Arc::new(Field::new(format!("state_{i}"), t, true))) + .collect(); + Self { + name, + accumulator, + signature, + return_type, + state_fields, + } + } + + /// Stored Python callable that returns a fresh accumulator instance + /// per partition. Consumed by the codec to cloudpickle the factory + /// across process boundaries. + pub(crate) fn accumulator(&self) -> &Py { + &self.accumulator + } + + pub(crate) fn return_type(&self) -> &DataType { + &self.return_type + } + + pub(crate) fn state_fields_ref(&self) -> &[FieldRef] { + &self.state_fields + } + + /// Reconstruct a `PythonFunctionAggregateUDF` from the parts emitted + /// by the codec. `state_fields` carries the full state schema + /// (names, data types, nullability, metadata) — the codec extracts + /// it from the IPC payload, so the post-decode state schema is + /// identical to the pre-encode one. Use [`Self::new`] when only + /// `Vec` is available (e.g. the Python constructor path, + /// where field names are synthesized). + pub(crate) fn from_parts( + name: String, + accumulator: Py, + input_types: Vec, + return_type: DataType, + state_fields: Vec, + volatility: Volatility, + ) -> Self { + Self { + name, + accumulator, + signature: Signature::exact(input_types, volatility), + return_type, + state_fields, + } + } +} + +impl Eq for PythonFunctionAggregateUDF {} +impl PartialEq for PythonFunctionAggregateUDF { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.signature == other.signature + && self.return_type == other.return_type + && self.state_fields == other.state_fields + // Pointer-identity fast path: `Arc`-shared clones of the + // same UDF skip the GIL roundtrip. Falls through to Python + // `__eq__` only for two distinct callables. + && (self.accumulator.as_ptr() == other.accumulator.as_ptr() + || Python::attach(|py| { + // See `PythonFunctionScalarUDF::eq` for the + // rationale on swallowing the exception as `false` + // and logging at `debug`. FIXME: revisit if + // upstream `AggregateUDFImpl` exposes a fallible + // `PartialEq`. + self.accumulator + .bind(py) + .eq(other.accumulator.bind(py)) + .unwrap_or_else(|e| { + log::debug!( + target: "datafusion_python::udaf", + "PythonFunctionAggregateUDF {:?} __eq__ raised; treating as unequal: {e}", + self.name, + ); + false + }) + })) + } +} + +impl std::hash::Hash for PythonFunctionAggregateUDF { + fn hash(&self, state: &mut H) { + // See `PythonFunctionScalarUDF`'s `Hash` impl for the + // rationale: hash the identifying header only and let + // `PartialEq` disambiguate callables. + self.name.hash(state); + self.signature.hash(state); + self.return_type.hash(state); + for f in &self.state_fields { + f.hash(state); + } + } +} + +impl AggregateUDFImpl for PythonFunctionAggregateUDF { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(self.return_type.clone()) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result> { + instantiate_accumulator(&self.accumulator) + } + + fn state_fields(&self, _args: StateFieldsArgs) -> Result> { + Ok(self.state_fields.clone()) + } } fn aggregate_udf_from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyDataFusionResult { @@ -190,14 +344,15 @@ impl PyAggregateUDF { state_type: PyArrowType>, volatility: &str, ) -> PyResult { - let function = create_udaf( - name, + let py_udf = PythonFunctionAggregateUDF::new( + name.to_string(), + accumulator, input_type.0, - Arc::new(return_type.0), + return_type.0, + state_type.0, parse_volatility(volatility)?, - to_rust_accumulator(accumulator), - Arc::new(state_type.0), ); + let function = AggregateUDF::new_from_impl(py_udf); Ok(Self { function }) } @@ -231,4 +386,9 @@ impl PyAggregateUDF { fn __repr__(&self) -> PyResult { Ok(format!("AggregateUDF({})", self.function.name())) } + + #[getter] + fn name(&self) -> &str { + self.function.name() + } } diff --git a/crates/core/src/udwf.rs b/crates/core/src/udwf.rs index 40e6208c4..ebec8f3bd 100644 --- a/crates/core/src/udwf.rs +++ b/crates/core/src/udwf.rs @@ -24,7 +24,6 @@ use datafusion::arrow::datatypes::DataType; use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}; use datafusion::error::{DataFusionError, Result}; use datafusion::logical_expr::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs}; -use datafusion::logical_expr::ptr_eq::PtrEq; use datafusion::logical_expr::window_state::WindowAggState; use datafusion::logical_expr::{ PartitionEvaluator, PartitionEvaluatorFactory, Signature, Volatility, WindowUDF, WindowUDFImpl, @@ -197,15 +196,24 @@ impl PartitionEvaluator for RustPartitionEvaluator { } } +fn instantiate_partition_evaluator(evaluator: &Py) -> Result> { + let instance = Python::attach(|py| { + evaluator + .call0(py) + .map_err(|e| DataFusionError::Execution(e.to_string())) + })?; + Ok(Box::new(RustPartitionEvaluator::new(instance))) +} + +/// Wrap a Python evaluator factory in a `PartitionEvaluatorFactory`. +/// +/// Retained for downstream callers that previously consumed this +/// helper to build a [`PartitionEvaluatorFactory`] for factory-based +/// APIs. New in-crate code should construct a +/// [`PythonFunctionWindowUDF`] directly so the codec can downcast and +/// ship it inline. pub fn to_rust_partition_evaluator(evaluator: Py) -> PartitionEvaluatorFactory { - Arc::new(move || -> Result> { - let evaluator = Python::attach(|py| { - evaluator - .call0(py) - .map_err(|e| DataFusionError::Execution(e.to_string())) - })?; - Ok(Box::new(RustPartitionEvaluator::new(evaluator))) - }) + Arc::new(move || instantiate_partition_evaluator(&evaluator)) } /// Represents an WindowUDF @@ -233,14 +241,14 @@ impl PyWindowUDF { volatility: &str, ) -> PyResult { let return_type = return_type.0; - let input_types = input_types.into_iter().map(|t| t.0).collect(); + let input_types: Vec = input_types.into_iter().map(|t| t.0).collect(); - let function = WindowUDF::from(MultiColumnWindowUDF::new( + let function = WindowUDF::from(PythonFunctionWindowUDF::new( name, + evaluator, input_types, return_type, parse_volatility(volatility)?, - to_rust_partition_evaluator(evaluator), )); Ok(Self { function }) } @@ -275,47 +283,104 @@ impl PyWindowUDF { fn __repr__(&self) -> PyResult { Ok(format!("WindowUDF({})", self.function.name())) } + + #[getter] + fn name(&self) -> &str { + self.function.name() + } } -#[derive(Hash, Eq, PartialEq)] -pub struct MultiColumnWindowUDF { +/// `WindowUDFImpl` for Python-defined window UDFs. +/// +/// Holds the Python evaluator factory directly so the codec can +/// downcast and cloudpickle it across process boundaries. Replaces +/// the prior factory-erased `MultiColumnWindowUDF`; the old name is +/// kept as a type alias below for backward compatibility. +#[derive(Debug)] +pub struct PythonFunctionWindowUDF { name: String, + evaluator: Py, signature: Signature, return_type: DataType, - partition_evaluator_factory: PtrEq, } -impl std::fmt::Debug for MultiColumnWindowUDF { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - f.debug_struct("WindowUDF") - .field("name", &self.name) - .field("signature", &self.signature) - .field("return_type", &"") - .field("partition_evaluator_factory", &"") - .finish() - } -} +/// Backward-compatible alias for downstream crates that referenced the +/// previous struct name. New code should use [`PythonFunctionWindowUDF`]. +pub type MultiColumnWindowUDF = PythonFunctionWindowUDF; -impl MultiColumnWindowUDF { +impl PythonFunctionWindowUDF { pub fn new( name: impl Into, + evaluator: Py, input_types: Vec, return_type: DataType, volatility: Volatility, - partition_evaluator_factory: PartitionEvaluatorFactory, ) -> Self { let name = name.into(); let signature = Signature::exact(input_types, volatility); Self { name, + evaluator, signature, return_type, - partition_evaluator_factory: partition_evaluator_factory.into(), } } + + /// Stored Python callable that produces a fresh partition + /// evaluator instance per partition. Consumed by the codec to + /// cloudpickle the evaluator factory across process boundaries. + pub(crate) fn evaluator(&self) -> &Py { + &self.evaluator + } + + pub(crate) fn return_type(&self) -> &DataType { + &self.return_type + } +} + +impl Eq for PythonFunctionWindowUDF {} +impl PartialEq for PythonFunctionWindowUDF { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.signature == other.signature + && self.return_type == other.return_type + // Pointer-identity fast path: `Arc`-shared clones of the + // same UDF skip the GIL roundtrip. Falls through to Python + // `__eq__` only for two distinct callables. + && (self.evaluator.as_ptr() == other.evaluator.as_ptr() + || Python::attach(|py| { + // See `PythonFunctionScalarUDF::eq` for the + // rationale on swallowing the exception as `false` + // and logging at `debug`. FIXME: revisit if + // upstream `WindowUDFImpl` exposes a fallible + // `PartialEq`. + self.evaluator + .bind(py) + .eq(other.evaluator.bind(py)) + .unwrap_or_else(|e| { + log::debug!( + target: "datafusion_python::udwf", + "PythonFunctionWindowUDF {:?} __eq__ raised; treating as unequal: {e}", + self.name, + ); + false + }) + })) + } +} + +impl std::hash::Hash for PythonFunctionWindowUDF { + fn hash(&self, state: &mut H) { + // See `PythonFunctionScalarUDF`'s `Hash` impl for the + // rationale: hash the identifying header only and let + // `PartialEq` disambiguate evaluators. + self.name.hash(state); + self.signature.hash(state); + self.return_type.hash(state); + } } -impl WindowUDFImpl for MultiColumnWindowUDF { +impl WindowUDFImpl for PythonFunctionWindowUDF { fn name(&self) -> &str { &self.name } @@ -334,7 +399,6 @@ impl WindowUDFImpl for MultiColumnWindowUDF { &self, _partition_evaluator_args: PartitionEvaluatorArgs, ) -> Result> { - let _ = _partition_evaluator_args; - (self.partition_evaluator_factory)() + instantiate_partition_evaluator(&self.evaluator) } } diff --git a/docs/source/user-guide/common-operations/aggregations.rst b/docs/source/user-guide/common-operations/aggregations.rst index f59b62ab4..8f218abd8 100644 --- a/docs/source/user-guide/common-operations/aggregations.rst +++ b/docs/source/user-guide/common-operations/aggregations.rst @@ -434,3 +434,21 @@ The available aggregate functions are: - :py:meth:`datafusion.expr.GroupingSet.cube` - :py:meth:`datafusion.expr.GroupingSet.grouping_sets` +User-Defined Aggregate Functions +-------------------------------- + +You can ship custom aggregations to the engine by subclassing +:py:class:`~datafusion.user_defined.Accumulator` and registering it via +:py:func:`~datafusion.udaf`. See :py:mod:`datafusion.user_defined` for +the accumulator interface and worked examples. + +.. note:: Serialization + + Python aggregate UDFs travel inline inside pickled or + :py:meth:`~datafusion.expr.Expr.to_bytes`-serialized expressions — + the accumulator class is captured by value via :mod:`cloudpickle`, + so worker processes do not need to pre-register the UDF. Any names + the accumulator resolves via ``import`` are captured **by reference** + and must be importable on the receiving worker. See + :py:mod:`datafusion.ipc` for the full IPC model and security caveats. + diff --git a/docs/source/user-guide/common-operations/windows.rst b/docs/source/user-guide/common-operations/windows.rst index d77881bcf..127f691b5 100644 --- a/docs/source/user-guide/common-operations/windows.rst +++ b/docs/source/user-guide/common-operations/windows.rst @@ -213,3 +213,21 @@ The possible window functions are: 3. Aggregate Functions - All :ref:`Aggregation Functions` can be used as window functions. + +User-Defined Window Functions +----------------------------- + +You can ship custom window functions to the engine by subclassing +:py:class:`~datafusion.user_defined.WindowEvaluator` and registering it +via :py:func:`~datafusion.udwf`. See :py:mod:`datafusion.user_defined` +for the evaluator interface and worked examples. + +.. note:: Serialization + + Python window UDFs travel inline inside pickled or + :py:meth:`~datafusion.expr.Expr.to_bytes`-serialized expressions — + the evaluator class is captured by value via :mod:`cloudpickle`, so + worker processes do not need to pre-register the UDF. Any names the + evaluator resolves via ``import`` are captured **by reference** and + must be importable on the receiving worker. See + :py:mod:`datafusion.ipc` for the full IPC model and security caveats. diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 645bd9c18..7e95bc127 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -449,14 +449,14 @@ def to_bytes(self, ctx: SessionContext | None = None) -> bytes: installed :class:`LogicalExtensionCodec`. When ``ctx`` is ``None``, the default codec is used. - Built-in functions and Python scalar UDFs travel inside the - returned bytes; the worker does not need to pre-register them. - UDFs imported via the FFI capsule protocol travel by name only - and must be registered on the worker. + Built-in functions and Python UDFs (scalar, aggregate, window) + travel inside the returned bytes; the worker does not need to + pre-register them. UDFs imported via the FFI capsule protocol + travel by name only and must be registered on the worker. .. warning:: Security Bytes returned here may embed a cloudpickled Python - callable (when the expression carries a Python scalar UDF). + callable (when the expression carries a Python UDF). Reconstructing them via :meth:`from_bytes` or :func:`pickle.loads` executes arbitrary Python on the receiver. Only accept payloads from trusted sources. @@ -526,7 +526,7 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr: ``ctx`` is ``None`` the worker context installed via :func:`datafusion.ipc.set_worker_ctx` is consulted; if no worker context is installed, the global :class:`SessionContext` is used - (sufficient for built-ins and Python scalar UDFs, plus any UDFs + (sufficient for built-ins and Python UDFs, plus any UDFs registered on the global context). .. warning:: Security @@ -560,12 +560,13 @@ def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]: Lets expressions be shipped to worker processes via :func:`pickle.dumps` / :func:`pickle.loads`. Built-in functions - and Python scalar UDFs travel inside the pickle bytes; only - FFI-capsule UDFs require pre-registration on the worker. The - worker's :class:`SessionContext` for resolving those references - is looked up via :func:`datafusion.ipc.set_worker_ctx`, falling - back to the global :class:`SessionContext` if none has been - installed on the worker. + and Python UDFs (scalar, aggregate, window) travel inside the + pickle bytes; only FFI-capsule UDFs require pre-registration on + the worker. The worker's :class:`SessionContext` for resolving + those references is looked up via + :func:`datafusion.ipc.set_worker_ctx`, falling back to the + global :class:`SessionContext` if none has been installed on + the worker. .. warning:: Security :func:`pickle.loads` on the returned tuple executes diff --git a/python/datafusion/ipc.py b/python/datafusion/ipc.py index 78b6873f7..8dd7fc463 100644 --- a/python/datafusion/ipc.py +++ b/python/datafusion/ipc.py @@ -35,16 +35,17 @@ def init_worker(): ctx.register_udaf(my_ffi_aggregate) set_worker_ctx(ctx) -Built-in functions and Python scalar UDFs travel inside the shipped -expression itself and do not need pre-registration on the worker. +Built-in functions and Python UDFs (scalar, aggregate, window) travel +inside the shipped expression itself and do not need pre-registration +on the worker. .. note:: Serialization model - Expressions containing Python scalar UDFs are serialized using - :mod:`cloudpickle`. The callable itself travels **by value** - (bytecode and closure cells inlined), but any names the callable - resolves via ``import`` are captured **by reference** and must be - importable on the receiving worker. + Expressions containing Python UDFs (scalar, aggregate, window) are + serialized using :mod:`cloudpickle`. The callable itself travels + **by value** (bytecode and closure cells inlined), but any names the + callable resolves via ``import`` are captured **by reference** and + must be importable on the receiving worker. The serialized payload is stamped with the sender's Python ``(major, minor)`` version. Loading on a different minor version @@ -97,8 +98,8 @@ def clear_worker_ctx() -> None: After clearing, expressions reconstructed in this worker fall back to the global :class:`SessionContext` — adequate for built-ins and Python - scalar UDFs, but anything imported via the FFI capsule protocol must - be registered on the global context to resolve. + UDFs (scalar, aggregate, window), but anything imported via the FFI + capsule protocol must be registered on the global context to resolve. Examples: >>> from datafusion import SessionContext diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index d79cf22e8..3eb50a094 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -441,6 +441,16 @@ def __init__( str(volatility), ) + @property + def name(self) -> str: + """Return the registered name of this UDAF. + + For UDAFs imported via the FFI capsule protocol, this is the + name the capsule itself reports — not the ``name`` argument + passed to the constructor (which is ignored on the FFI path). + """ + return self._udaf.name + def __repr__(self) -> str: """Print a string representation of the Aggregate UDF.""" return self._udaf.__repr__() @@ -851,6 +861,16 @@ def __init__( name, func, input_types, return_type, str(volatility) ) + @property + def name(self) -> str: + """Return the registered name of this UDWF. + + For UDWFs imported via the FFI capsule protocol, this is the + name the capsule itself reports — not the ``name`` argument + passed to the constructor (which is ignored on the FFI path). + """ + return self._udwf.name + def __repr__(self) -> str: """Print a string representation of the Window UDF.""" return self._udwf.__repr__() diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py index 5d8d9285f..eb0441c49 100644 --- a/python/tests/test_pickle_expr.py +++ b/python/tests/test_pickle_expr.py @@ -17,10 +17,10 @@ """In-process pickle round-trip tests for :class:`Expr`. -Built-in functions and Python scalar UDFs travel with the pickled -expression and do not need worker-side pre-registration. The worker -context (:mod:`datafusion.ipc`) is only consulted for UDFs imported -via the FFI capsule protocol. +Built-in functions and Python UDFs (scalar, aggregate, window) travel +with the pickled expression and do not need worker-side pre-registration. +The worker context (:mod:`datafusion.ipc`) is only consulted for UDFs +imported via the FFI capsule protocol. """ from __future__ import annotations @@ -147,6 +147,135 @@ def test_multi_arg_udf_round_trip(self): assert "add_scaled" in decoded.canonical_name() +class TestAggregateUDFCodec: + """Python aggregate UDFs travel inline like scalar UDFs.""" + + def _build_aggregate_udf(self): + from datafusion import udaf + from datafusion.user_defined import Accumulator + + class CountAcc(Accumulator): + def __init__(self): + self._count = 0 + + def state(self): + return [pa.scalar(self._count, type=pa.int64())] + + def update(self, values): + self._count += len(values) + + def merge(self, states): + partition_counts = states[0] + for i in range(len(partition_counts)): + self._count += partition_counts[i].as_py() + + def evaluate(self): + return pa.scalar(self._count, type=pa.int64()) + + return udaf( + CountAcc, + [pa.int64()], + pa.int64(), + [pa.int64()], + "immutable", + name="count_all", + ) + + def test_agg_udf_self_contained_blob(self): + u = self._build_aggregate_udf() + e = u(col("a")) + blob = pickle.dumps(e) + assert len(blob) > 200 + + def test_agg_udf_decodes_into_fresh_ctx(self): + u = self._build_aggregate_udf() + e = u(col("a")) + blob = e.to_bytes() + fresh = SessionContext() + decoded = Expr.from_bytes(blob, ctx=fresh) + assert "count_all" in decoded.canonical_name() + + def test_agg_udf_decodes_via_pickle_with_no_worker_ctx(self): + u = self._build_aggregate_udf() + e = u(col("a")) + blob = pickle.dumps(e) + decoded = pickle.loads(blob) # noqa: S301 + assert "count_all" in decoded.canonical_name() + + def test_agg_udf_evaluates_after_roundtrip(self): + """End-to-end: the decoded aggregate UDF runs and merges across + partitions, exercising the round-tripped state-field schema.""" + u = self._build_aggregate_udf() + e = u(col("a")) + decoded = pickle.loads(pickle.dumps(e)) # noqa: S301 + + ctx = SessionContext() + schema = pa.schema([pa.field("a", pa.int64())]) + batch1 = pa.record_batch([pa.array([1, 2, 3], type=pa.int64())], schema=schema) + batch2 = pa.record_batch([pa.array([4, 5], type=pa.int64())], schema=schema) + df = ctx.create_dataframe([[batch1], [batch2]]) + out = df.aggregate([], [decoded.alias("n")]).to_pydict() + assert out["n"] == [5] + + +class TestWindowUDFCodec: + """Python window UDFs travel inline like scalar UDFs.""" + + def _build_window_udf(self): + from datafusion import udwf + from datafusion.user_defined import WindowEvaluator + + class CountUpEvaluator(WindowEvaluator): + def evaluate_all(self, values, num_rows): + return pa.array(list(range(num_rows))) + + return udwf( + CountUpEvaluator, + [pa.int64()], + pa.int64(), + "immutable", + name="count_up", + ) + + def test_window_udf_self_contained_blob(self): + u = self._build_window_udf() + e = u(col("a")) + blob = pickle.dumps(e) + assert len(blob) > 200 + + def test_window_udf_decodes_into_fresh_ctx(self): + u = self._build_window_udf() + e = u(col("a")) + blob = e.to_bytes() + fresh = SessionContext() + decoded = Expr.from_bytes(blob, ctx=fresh) + assert "count_up" in decoded.canonical_name() + + def test_window_udf_decodes_via_pickle_with_no_worker_ctx(self): + u = self._build_window_udf() + e = u(col("a")) + blob = pickle.dumps(e) + decoded = pickle.loads(blob) # noqa: S301 + assert "count_up" in decoded.canonical_name() + + def test_window_udf_evaluates_after_roundtrip(self): + """End-to-end: decoded window UDF runs and emits per-row values + produced by the round-tripped evaluator factory.""" + from datafusion.expr import WindowFrame + + u = self._build_window_udf() + e = u(col("a")) + decoded = pickle.loads(pickle.dumps(e)) # noqa: S301 + + ctx = SessionContext() + df = ctx.from_pydict({"a": [1, 2, 3, 4, 5]}) + framed = ( + decoded.window_frame(WindowFrame("rows", None, None)).build().alias("c") + ) + out = df.select(framed).to_pydict() + assert out["c"] == [0, 1, 2, 3, 4] + + class TestErrorPaths: def test_from_bytes_rejects_garbage(self): with pytest.raises(Exception): # noqa: B017 From f43830480b754900b93a30c36e6280d0ce0577f1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 26 May 2026 11:48:57 -0400 Subject: [PATCH 38/83] feat: expose variety of features from DF54 update (#1554) * refactor: migrate FFI example table function to call_with_args DataFusion 53 deprecated `TableFunctionImpl::call(args: &[Expr])` in favor of `call_with_args(args: TableFunctionArgs)`. `PyTableFunction` was migrated in 5a64b0d; this brings the FFI example along so it no longer relies on the deprecated entry point. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: type SessionContext codec setters with exportable Protocols PR #1541 introduced `with_logical_extension_codec` / `with_physical_extension_codec` setters typed as `codec: Any`. The Rust extractors accept either a raw `PyCapsule` or any object exposing `__datafusion_logical_extension_codec__` / `__datafusion_physical_extension_codec__`. Add `LogicalExtensionCodecExportable` / `PhysicalExtensionCodecExportable` Protocols in `python/datafusion/user_defined.py` (matching the existing `ScalarUDFExportable` pattern) and tighten both setter signatures to `Protocol | _PyCapsule`. Pure typing change; no runtime behavior diff. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: accept variadic field path in get_field Upstream exposes both `get_field(expr, name)` and `get_field_path(expr, [names...])`, but both ultimately call the same scalar UDF with a base expression plus one or more name args. Collapse the Python surface into a single variadic `get_field(expr, *names)` that accepts either a one-step lookup or a path of names, dispatching through a single Rust binding. Note in `.ai/skills/check-upstream/SKILL.md` that `get_field_path` is covered by the variadic form so future audits do not flag it as a gap. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: SessionContext.read_batches / read_batch Wrap upstream `SessionContext::read_batches`, which materializes a DataFrame directly from a sequence of `RecordBatch`es without registering a named table. The single-batch convenience `SessionContext.read_batch` is implemented in pure Python by calling `read_batches([batch])`, so the Rust side only needs the one binding. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: SessionContext UDF lookup helpers Expose `udf(name)` / `udaf(name)` / `udwf(name)` lookups symmetric with the existing `register_udf` / `register_udaf` / `register_udwf` setters, plus `udfs()` / `udafs()` / `udwfs()` for enumerating registered function names. Looked-up functions come back as the same `ScalarUDF` / `AggregateUDF` / `WindowUDF` wrappers users already get from registration, so they can be called as expressions or re-registered into a different session. Returns Vec from the list helpers (sorted) rather than the raw HashSet upstream returns, so calling code gets a stable ordering. Co-Authored-By: Claude Opus 4.7 (1M context) * bump pre-commit so it stops failing CI checks * test: drop xfail on timestamp[s] parquet roundtrip pyarrow.parquet promotes timestamp[s] to timestamp[ms] on write (apache/arrow#41382), so the read array never matched the input. Cast the expected array to timestamp[ms] in test_simple_select to assert DataFusion reads what Arrow actually stored. Co-Authored-By: Claude Opus 4.7 (1M context) * test: capture deprecation warning in repr_rows conflict case DataFrameHtmlFormatter(repr_rows=..., max_rows=...) fires the deprecation warning before raising ValueError, but pytest.raises does not catch warnings. The escaping warning surfaced in every pytest run. Wrap the call in both pytest.raises and pytest.warns so the warning is asserted, not leaked. Co-Authored-By: Claude Opus 4.7 (1M context) * docs(udf): document SessionContext UDF lookup with worked examples Add Examples docstrings (doctest) for `udf` / `udaf` / `udwf` / `udfs` / `udafs` / `udwfs` that demonstrate the lookup pattern, including a late-binding example where the function name comes from configuration. Add tests covering config-driven dispatch and built-in UDAF / UDWF lookup so the documented patterns are exercised end-to-end. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor(udf): raise KeyError on UDF/UDAF/UDWF lookup miss `SessionContext.udf` / `udaf` / `udwf` previously surfaced upstream `DataFusionError::Plan` as a generic exception whose message ("There is no UDF named ...") is set by DataFusion and can drift between releases. Pre-check membership via `udfs()` / `udafs()` / `udwfs()` and raise `PyKeyError` on miss so callers get the Pythonic dict-style lookup behavior and tests are no longer coupled to the upstream wording. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor(udf): add _from_internal classmethod to UDF wrappers `SessionContext.udf` / `udaf` / `udwf` previously constructed wrapper objects by calling `__new__` directly and writing the private `_udf` / `_udaf` / `_udwf` attribute from outside the owning module. Three near-identical blocks coupled `context.py` to wrapper internals. Add a `_from_internal` classmethod on each wrapper that takes an already-constructed `df_internal` handle and returns a wrapper without re-running `__init__`. The lookup methods now collapse to a single call, the `__new__` bypass is documented on the wrapper class itself, and renaming the private field is a one-spot edit. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: widen SessionContext.read_batches to accept any iterable The underlying PyArrow FFI extractor for `Vec` requires a Python `list`, so the previous `list[pa.RecordBatch]` annotation was accurate but unnecessarily strict. Accept any `Iterable[pa.RecordBatch]` on the Python side and materialize to a list before crossing the FFI boundary so callers can pass generators, tuples, or other iterables without manual conversion. Co-Authored-By: Claude Opus 4.7 (1M context) * docs(context): trim codec docstrings, reference Exportable protocols Drop prose restatement of the type union for `with_logical_extension_codec` and `with_physical_extension_codec`. Keep the dunder name (not visible from the type hint) and cross-link the `LogicalExtensionCodecExportable` / `PhysicalExtensionCodecExportable` protocols so Sphinx resolves them. Co-Authored-By: Claude Opus 4.7 (1M context) * docs(udf): drop return-type cross-refs in udf/udaf/udwf docstrings The `:py:class:` link back to the wrapper class shadowed the return type annotation and risked drifting if the class were moved. Replace with a plain backtick literal; surrounding contract prose is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) * docs(functions): use F alias in get_field doctest The doctest namespace already imports `datafusion.functions as F`, making `F.named_struct` / `F.get_field` shorter than the `dfn.functions.*` form. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .ai/skills/check-upstream/SKILL.md | 8 +- .pre-commit-config.yaml | 2 +- crates/core/src/context.rs | 51 ++++- crates/core/src/functions.rs | 8 +- .../src/table_function.rs | 5 +- python/datafusion/context.py | 208 +++++++++++++++++- python/datafusion/functions.py | 46 +++- python/datafusion/user_defined.py | 48 ++++ python/tests/test_context.py | 26 +++ python/tests/test_dataframe.py | 8 +- python/tests/test_functions.py | 31 +++ python/tests/test_sql.py | 14 +- python/tests/test_udf.py | 71 ++++++ 13 files changed, 491 insertions(+), 35 deletions(-) diff --git a/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md index 3bac018ef..23873feab 100644 --- a/.ai/skills/check-upstream/SKILL.md +++ b/.ai/skills/check-upstream/SKILL.md @@ -66,11 +66,17 @@ The user may specify an area via `$ARGUMENTS`. If no area is specified or "all" - Python API: `python/datafusion/functions.py` — each function wraps a call to `datafusion._internal.functions` - Rust bindings: `crates/core/src/functions.rs` — `#[pyfunction]` definitions registered via `init_module()` +**Evaluated and not requiring separate Python exposure:** +- `get_field_path` — already covered by `get_field(expr, *names)`, which takes a + variadic field path and dispatches to the same underlying + `functions::core::get_field` UDF as the upstream `get_field_path` helper. + **How to check:** 1. Fetch the upstream scalar function documentation page 2. Compare against functions listed in `python/datafusion/functions.py` (check the `__all__` list and function definitions) 3. A function is covered if it exists in the Python API — it does NOT need a dedicated Rust `#[pyfunction]`. Many functions are aliases that reuse another function's Rust binding. -4. Only report functions that are missing from the Python `__all__` list / function definitions +4. Check against the "evaluated and not requiring exposure" list before flagging as a gap +5. Only report functions that are missing from the Python `__all__` list / function definitions ### 2. Aggregate Functions diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2d3c2bc59..0a212480b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,7 +17,7 @@ repos: - repo: https://github.com/rhysd/actionlint - rev: v1.7.6 + rev: v1.7.12 hooks: - id: actionlint-docker - repo: https://github.com/astral-sh/ruff-pre-commit diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 642afeef7..bca8e30f5 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -35,7 +35,6 @@ use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; use datafusion::datasource::{MemTable, TableProvider}; -use datafusion::execution::TaskContextProvider; use datafusion::execution::context::{ DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext, }; @@ -44,6 +43,7 @@ use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, Unboun use datafusion::execution::options::{ArrowReadOptions, ReadOptions}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::session_state::SessionStateBuilder; +use datafusion::execution::{FunctionRegistry, TaskContextProvider}; use datafusion::prelude::{ AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions, }; @@ -847,6 +847,13 @@ impl PySessionContext { Ok(()) } + pub fn read_batches( + &self, + batches: PyArrowType>, + ) -> PyDataFusionResult { + Ok(PyDataFrame::new(self.ctx.read_batches(batches.0)?)) + } + #[allow(clippy::too_many_arguments)] #[pyo3(signature = (name, path, table_partition_cols=vec![], parquet_pruning=true, @@ -1065,6 +1072,48 @@ impl PySessionContext { self.ctx.deregister_udwf(name); } + pub fn udf(&self, name: &str) -> PyResult { + if !self.ctx.udfs().contains(name) { + return Err(PyKeyError::new_err(format!("no UDF named '{name}'"))); + } + let function = (*self.ctx.udf(name).map_err(py_datafusion_err)?).clone(); + Ok(PyScalarUDF { function }) + } + + pub fn udaf(&self, name: &str) -> PyResult { + if !self.ctx.udafs().contains(name) { + return Err(PyKeyError::new_err(format!("no UDAF named '{name}'"))); + } + let function = (*self.ctx.udaf(name).map_err(py_datafusion_err)?).clone(); + Ok(PyAggregateUDF { function }) + } + + pub fn udwf(&self, name: &str) -> PyResult { + if !self.ctx.udwfs().contains(name) { + return Err(PyKeyError::new_err(format!("no UDWF named '{name}'"))); + } + let function = (*self.ctx.udwf(name).map_err(py_datafusion_err)?).clone(); + Ok(PyWindowUDF { function }) + } + + pub fn udfs(&self) -> Vec { + let mut names: Vec = self.ctx.udfs().into_iter().collect(); + names.sort(); + names + } + + pub fn udafs(&self) -> Vec { + let mut names: Vec = self.ctx.udafs().into_iter().collect(); + names.sort(); + names + } + + pub fn udwfs(&self) -> Vec { + let mut names: Vec = self.ctx.udwfs().into_iter().collect(); + names.sort(); + names + } + #[pyo3(signature = (name="datafusion"))] pub fn catalog(&self, py: Python, name: &str) -> PyResult> { let catalog = self.ctx.catalog(name).ok_or(PyKeyError::new_err(format!( diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 7feb62d79..5f47d123b 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -574,10 +574,10 @@ expr_fn!(union_tag, arg1); expr_fn!(random); #[pyfunction] -fn get_field(expr: PyExpr, name: PyExpr) -> PyExpr { - functions::core::get_field() - .call(vec![expr.into(), name.into()]) - .into() +fn get_field(expr: PyExpr, names: Vec) -> PyExpr { + let mut args = vec![expr.into()]; + args.extend(names.into_iter().map(Into::into)); + functions::core::get_field().call(args).into() } #[pyfunction] diff --git a/examples/datafusion-ffi-example/src/table_function.rs b/examples/datafusion-ffi-example/src/table_function.rs index 79c13f64d..ed3ef142b 100644 --- a/examples/datafusion-ffi-example/src/table_function.rs +++ b/examples/datafusion-ffi-example/src/table_function.rs @@ -17,9 +17,8 @@ use std::sync::Arc; -use datafusion_catalog::{TableFunctionImpl, TableProvider}; +use datafusion_catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider}; use datafusion_common::error::Result as DataFusionResult; -use datafusion_expr::Expr; use datafusion_ffi::udtf::FFI_TableFunction; use datafusion_python_util::ffi_logical_codec_from_pycapsule; use pyo3::types::PyCapsule; @@ -59,7 +58,7 @@ impl MyTableFunction { } impl TableFunctionImpl for MyTableFunction { - fn call(&self, _args: &[Expr]) -> DataFusionResult> { + fn call_with_args(&self, _args: TableFunctionArgs) -> DataFusionResult> { let provider = MyTableProvider::new(4, 3, 2).create_table()?; Ok(Arc::new(provider)) } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 5c3501941..9ecbfe311 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -82,10 +82,11 @@ if TYPE_CHECKING: import pathlib - from collections.abc import Sequence + from collections.abc import Iterable, Sequence import pandas as pd import polars as pl # type: ignore[import] + from _typeshed import CapsuleType as _PyCapsule from datafusion.catalog import CatalogProvider, Table from datafusion.common import DFSchema @@ -93,6 +94,8 @@ from datafusion.plan import ExecutionPlan, LogicalPlan from datafusion.user_defined import ( AggregateUDF, + LogicalExtensionCodecExportable, + PhysicalExtensionCodecExportable, ScalarUDF, TableFunction, WindowUDF, @@ -959,6 +962,52 @@ def register_record_batches( """ self.ctx.register_record_batches(name, partitions) + def read_batch(self, batch: pa.RecordBatch) -> DataFrame: + """Return a :py:class:`~datafusion.DataFrame` reading a single batch. + + Convenience wrapper around :py:meth:`read_batches` for the single-batch + case. Unlike :py:meth:`register_batch`, this does not register the + batch as a named table; it returns an anonymous + :py:class:`~datafusion.DataFrame` directly. + + Args: + batch: Record batch to wrap as a DataFrame. + + Examples: + >>> ctx = dfn.SessionContext() + >>> batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3]}) + >>> ctx.read_batch(batch).to_pydict() + {'a': [1, 2, 3]} + """ + return self.read_batches([batch]) + + def read_batches(self, batches: Iterable[pa.RecordBatch]) -> DataFrame: + """Return a :py:class:`~datafusion.DataFrame` reading the given batches. + + All batches must share the same schema. Any iterable of + :py:class:`pa.RecordBatch` is accepted (list, tuple, generator); + it is materialized into a list before being handed to the + underlying Rust binding. Unlike :py:meth:`register_record_batches`, + this does not register the batches as a named table; it returns + an anonymous :py:class:`~datafusion.DataFrame` directly. + + Args: + batches: Record batches to wrap as a DataFrame. + + Examples: + >>> ctx = dfn.SessionContext() + >>> b1 = pa.RecordBatch.from_pydict({"a": [1, 2]}) + >>> b2 = pa.RecordBatch.from_pydict({"a": [3, 4]}) + >>> ctx.read_batches([b1, b2]).to_pydict() + {'a': [1, 2, 3, 4]} + + A generator works too: + + >>> ctx.read_batches(b for b in [b1, b2]).to_pydict() + {'a': [1, 2, 3, 4]} + """ + return DataFrame(self.ctx.read_batches(list(batches))) + def register_parquet( self, name: str, @@ -1268,6 +1317,145 @@ def deregister_udwf(self, name: str) -> None: """ self.ctx.deregister_udwf(name) + def udf(self, name: str) -> ScalarUDF: + """Look up a registered scalar UDF by name. + + Returns the same ``ScalarUDF`` wrapper that :py:meth:`register_udf` + accepts, so it can be invoked as an expression in the DataFrame API + or re-registered into a different :py:class:`SessionContext`. + Built-in scalar functions from the session's function registry are + also looked up. + + Args: + name: Name of the registered scalar UDF. + + Raises: + KeyError: If no scalar UDF is registered under ``name``. + + Examples: + Register a UDF, then look it up by name and use it in the + DataFrame API: + + >>> ctx = dfn.SessionContext() + >>> nullcheck = dfn.udf( + ... lambda x: x.is_null(), + ... [pa.int64()], + ... pa.bool_(), + ... volatility="immutable", + ... name="nullcheck", + ... ) + >>> ctx.register_udf(nullcheck) + >>> fn = ctx.udf("nullcheck") + >>> df = ctx.from_pydict({"a": [1, None, 3]}) + >>> df.select(fn(col("a")).alias("is_null")).to_pydict() + {'is_null': [False, True, False]} + + Late-binding: the function name can come from configuration + rather than an imported symbol, which is useful when the set + of UDFs is plugin-driven or chosen at runtime: + + >>> config = {"null_check": "nullcheck"} + >>> fn = ctx.udf(config["null_check"]) + >>> df.select(fn(col("a")).alias("is_null")).to_pydict() + {'is_null': [False, True, False]} + """ + from datafusion.user_defined import ScalarUDF as _ScalarUDF # noqa: PLC0415 + + return _ScalarUDF._from_internal(self.ctx.udf(name)) + + def udaf(self, name: str) -> AggregateUDF: + """Look up a registered aggregate UDF by name. + + Returns the same ``AggregateUDF`` wrapper that :py:meth:`register_udaf` + accepts. Built-in aggregate functions such as ``sum`` or ``avg`` are + also discoverable through this lookup. See :py:meth:`udf` for a worked + late-binding example; the pattern is identical for aggregates. + + Args: + name: Name of the registered aggregate UDF. + + Raises: + KeyError: If no aggregate UDF is registered under ``name``. + + Examples: + Look up a built-in aggregate by name and use it in + :py:meth:`~datafusion.DataFrame.aggregate`: + + >>> ctx = dfn.SessionContext() + >>> sum_fn = ctx.udaf("sum") + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> df.aggregate([], [sum_fn(col("a")).alias("total")]).to_pydict() + {'total': [6]} + """ + from datafusion.user_defined import ( # noqa: PLC0415 + AggregateUDF as _AggregateUDF, + ) + + return _AggregateUDF._from_internal(self.ctx.udaf(name)) + + def udwf(self, name: str) -> WindowUDF: + """Look up a registered window UDF by name. + + Returns the same ``WindowUDF`` wrapper that :py:meth:`register_udwf` + accepts. Built-in window functions such as ``row_number`` or ``rank`` + are also discoverable through this lookup. See :py:meth:`udf` for a + worked late-binding example; the pattern is identical for window + functions. + + Args: + name: Name of the registered window UDF. + + Raises: + KeyError: If no window UDF is registered under ``name``. + + Examples: + Look up a built-in window function by name and use it in + ``select``: + + >>> ctx = dfn.SessionContext() + >>> rn = ctx.udwf("row_number") + >>> df = ctx.from_pydict({"a": [10, 20, 30]}) + >>> df.select(col("a"), rn().alias("rn")).to_pydict() + {'a': [10, 20, 30], 'rn': [1, 2, 3]} + """ + from datafusion.user_defined import WindowUDF as _WindowUDF # noqa: PLC0415 + + return _WindowUDF._from_internal(self.ctx.udwf(name)) + + def udfs(self) -> list[str]: + """Return the sorted names of all registered scalar UDFs. + + Includes both user-registered and built-in scalar functions. Pair + with :py:meth:`udf` to drive discovery, validation, or config-based + dispatch. + + Examples: + >>> ctx = dfn.SessionContext() + >>> "abs" in ctx.udfs() + True + """ + return self.ctx.udfs() + + def udafs(self) -> list[str]: + """Return the sorted names of all registered aggregate UDFs. + + Examples: + >>> ctx = dfn.SessionContext() + >>> "sum" in ctx.udafs() + True + """ + return self.ctx.udafs() + + def udwfs(self) -> list[str]: + """Return the sorted names of all registered window UDFs. + + Examples: + >>> ctx = dfn.SessionContext() + >>> "row_number" in ctx.udwfs() + True + """ + return self.ctx.udwfs() + def catalog(self, name: str = "datafusion") -> Catalog: """Retrieve a catalog by name.""" return Catalog(self.ctx.catalog(name)) @@ -1744,11 +1932,14 @@ def __datafusion_logical_extension_codec__(self) -> Any: """Access the PyCapsule FFI_LogicalExtensionCodec.""" return self.ctx.__datafusion_logical_extension_codec__() - def with_logical_extension_codec(self, codec: Any) -> SessionContext: + def with_logical_extension_codec( + self, codec: LogicalExtensionCodecExportable | _PyCapsule + ) -> SessionContext: """Create a new session context with specified codec. - This only supports codecs that have been implemented using the - FFI interface. + Only FFI codecs are supported. Pass any object implementing + ``__datafusion_logical_extension_codec__`` (see + :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`). """ new_internal = self.ctx.with_logical_extension_codec(codec) new = SessionContext.__new__(SessionContext) @@ -1759,11 +1950,14 @@ def __datafusion_physical_extension_codec__(self) -> Any: """Access the PyCapsule FFI_PhysicalExtensionCodec.""" return self.ctx.__datafusion_physical_extension_codec__() - def with_physical_extension_codec(self, codec: Any) -> SessionContext: + def with_physical_extension_codec( + self, codec: PhysicalExtensionCodecExportable | _PyCapsule + ) -> SessionContext: """Create a new session context with the specified physical codec. - This only supports codecs that have been implemented using the - FFI interface. + Only FFI codecs are supported. Pass any object implementing + ``__datafusion_physical_extension_codec__`` (see + :py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`). """ new_internal = self.ctx.with_physical_extension_codec(codec) new = SessionContext.__new__(SessionContext) diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 9761d1879..28c10e005 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -2727,24 +2727,32 @@ def arrow_metadata(expr: Expr, key: Expr | str | None = None) -> Expr: return Expr(f.arrow_metadata(expr.expr, key.expr)) -def get_field(expr: Expr, name: Expr | str) -> Expr: - """Extracts a field from a struct or map by name. +def get_field(expr: Expr, *names: Expr | str) -> Expr: + """Extracts a (possibly nested) field from a struct or map by name. - When the field name is a static string, the bracket operator - ``expr["field"]`` is a convenient shorthand. Use ``get_field`` - when the field name is a dynamic expression. + Pass one name for a single-level lookup, or several names to walk a path + of nested struct/map fields in a single ``get_field`` call. For a single + static-string name, ``expr["field"]`` is a convenient shorthand; use + ``get_field`` when the field name is a dynamic + :py:class:`~datafusion.expr.Expr` or when traversing multiple levels at + once. + + Args: + expr: The struct or map expression to read from. + *names: One or more field names (``str``) or expressions + (:py:class:`~datafusion.expr.Expr`). Examples: + Single-level lookup: + >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [1], "b": [2]}) >>> df = df.with_column( ... "s", - ... dfn.functions.named_struct( - ... [("x", dfn.col("a")), ("y", dfn.col("b"))] - ... ), + ... F.named_struct([("x", dfn.col("a")), ("y", dfn.col("b"))]), ... ) >>> result = df.select( - ... dfn.functions.get_field(dfn.col("s"), "x").alias("x_val") + ... F.get_field(dfn.col("s"), "x").alias("x_val") ... ) >>> result.collect_column("x_val")[0].as_py() 1 @@ -2756,10 +2764,24 @@ def get_field(expr: Expr, name: Expr | str) -> Expr: ... ) >>> result.collect_column("x_val")[0].as_py() 1 + + Multi-level lookup: + + >>> df = df.with_column( + ... "outer", + ... F.named_struct([("inner", dfn.col("s"))]), + ... ) + >>> result = df.select( + ... F.get_field(dfn.col("outer"), "inner", "x").alias("x_val") + ... ) + >>> result.collect_column("x_val")[0].as_py() + 1 """ - if isinstance(name, str): - name = Expr.string_literal(name) - return Expr(f.get_field(expr.expr, name.expr)) + if not names: + msg = "get_field requires at least one field name" + raise ValueError(msg) + resolved = [Expr.string_literal(n) if isinstance(n, str) else n for n in names] + return Expr(f.get_field(expr.expr, [n.expr for n in resolved])) def union_extract(union_expr: Expr, field_name: Expr | str) -> Expr: diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index 3eb50a094..a8ee7756e 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -113,6 +113,18 @@ def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]: return value.__class__.__name__ == "PyCapsule" +class LogicalExtensionCodecExportable(Protocol): + """Type hint for objects exposing ``__datafusion_logical_extension_codec__``.""" + + def __datafusion_logical_extension_codec__(self) -> object: ... # noqa: D105 + + +class PhysicalExtensionCodecExportable(Protocol): + """Type hint for objects exposing ``__datafusion_physical_extension_codec__``.""" + + def __datafusion_physical_extension_codec__(self) -> object: ... # noqa: D105 + + class ScalarUDF: """Class for performing scalar user-defined functions (UDF). @@ -141,6 +153,18 @@ def __init__( name, func, input_fields, return_field, str(volatility) ) + @classmethod + def _from_internal(cls, internal: df_internal.ScalarUDF) -> ScalarUDF: + """Wrap an already-constructed internal ``ScalarUDF`` handle. + + Used by :py:meth:`SessionContext.udf` to surface a function looked + up from the session's function registry without re-running + :py:meth:`__init__`. + """ + wrapper = cls.__new__(cls) + wrapper._udf = internal + return wrapper + @property def name(self) -> str: """Return the registered name of this UDF. @@ -441,6 +465,18 @@ def __init__( str(volatility), ) + @classmethod + def _from_internal(cls, internal: df_internal.AggregateUDF) -> AggregateUDF: + """Wrap an already-constructed internal ``AggregateUDF`` handle. + + Used by :py:meth:`SessionContext.udaf` to surface a function looked + up from the session's function registry without re-running + :py:meth:`__init__`. + """ + wrapper = cls.__new__(cls) + wrapper._udaf = internal + return wrapper + @property def name(self) -> str: """Return the registered name of this UDAF. @@ -861,6 +897,18 @@ def __init__( name, func, input_types, return_type, str(volatility) ) + @classmethod + def _from_internal(cls, internal: df_internal.WindowUDF) -> WindowUDF: + """Wrap an already-constructed internal ``WindowUDF`` handle. + + Used by :py:meth:`SessionContext.udwf` to surface a function looked + up from the session's function registry without re-running + :py:meth:`__init__`. + """ + wrapper = cls.__new__(cls) + wrapper._udwf = internal + return wrapper + @property def name(self) -> str: """Return the registered name of this UDWF. diff --git a/python/tests/test_context.py b/python/tests/test_context.py index e0ebdbae5..112a6fd7b 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -905,6 +905,32 @@ def test_register_batch_empty(ctx): assert result[0].num_rows == 0 +def test_read_batch_returns_dataframe(ctx): + batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}) + df = ctx.read_batch(batch) + assert df.to_pydict() == {"a": [1, 2, 3], "b": [4, 5, 6]} + # read_batch should not register a named table. + assert ctx.catalog().schema().names() == set() + + +def test_read_batches_concatenates(ctx): + b1 = pa.RecordBatch.from_pydict({"a": [1, 2]}) + b2 = pa.RecordBatch.from_pydict({"a": [3, 4]}) + df = ctx.read_batches([b1, b2]) + assert df.to_pydict() == {"a": [1, 2, 3, 4]} + + +def test_read_batches_accepts_iterable(ctx): + b1 = pa.RecordBatch.from_pydict({"a": [1, 2]}) + b2 = pa.RecordBatch.from_pydict({"a": [3, 4]}) + # Generator: ensures non-list iterables are materialized before FFI. + df = ctx.read_batches(b for b in (b1, b2)) + assert df.to_pydict() == {"a": [1, 2, 3, 4]} + # Tuple: same. + df = ctx.read_batches((b1, b2)) + assert df.to_pydict() == {"a": [1, 2, 3, 4]} + + def test_create_sql_options(): SQLOptions() diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index 6bd0ce9f9..ab3992a79 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -1704,8 +1704,12 @@ def test_repr_rows_backward_compatibility(clean_formatter_state): assert formatter.max_rows == 15 assert formatter.repr_rows == 15 - # Should fail when conflicting with max_rows - with pytest.raises(ValueError, match="Cannot specify both repr_rows and max_rows"): + # Should fail when conflicting with max_rows. The deprecation warning still + # fires before the ValueError, so assert both. + with ( + pytest.raises(ValueError, match="Cannot specify both repr_rows and max_rows"), + pytest.warns(DeprecationWarning, match="repr_rows parameter is deprecated"), + ): DataFrameHtmlFormatter(repr_rows=5, max_rows=10) # Setting repr_rows via property should warn diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 5538fc33b..55d9c8ee8 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -1957,6 +1957,37 @@ def test_get_field(df): assert result.column(1) == pa.array([4, 5, 6]) +def test_get_field_path(df): + df = df.with_column( + "outer", + f.named_struct( + [ + ( + "inner", + f.named_struct( + [ + ("x", column("a")), + ("y", column("b")), + ] + ), + ), + ] + ), + ) + result = df.select( + f.get_field(column("outer"), "inner", "x").alias("x_val"), + f.get_field(column("outer"), "inner", "y").alias("y_val"), + ).collect()[0] + + assert result.column(0) == pa.array(["Hello", "World", "!"], type=pa.string_view()) + assert result.column(1) == pa.array([4, 5, 6]) + + +def test_get_field_requires_a_name(): + with pytest.raises(ValueError, match="at least one field name"): + f.get_field(column("s")) + + def test_arrow_metadata(): ctx = SessionContext() field = pa.field("val", pa.int64(), metadata={"key1": "value1", "key2": "value2"}) diff --git a/python/tests/test_sql.py b/python/tests/test_sql.py index 1ed1746e1..924d2655c 100644 --- a/python/tests/test_sql.py +++ b/python/tests/test_sql.py @@ -450,13 +450,9 @@ def test_udf( pa.array([b"1111", b"2222", b"3333"], pa.binary(4), _null_mask), id="binary4", ), - # `timestamp[s]` does not roundtrip for pyarrow.parquet: https://github.com/apache/arrow/issues/41382 pytest.param( helpers.data_datetime("s"), id="datetime_s", - marks=pytest.mark.xfail( - reason="pyarrow.parquet does not support timestamp[s] roundtrips" - ), ), pytest.param( helpers.data_datetime("ms"), @@ -484,6 +480,16 @@ def test_simple_select(ctx, tmp_path, arr): batches = ctx.sql("SELECT a AS tt FROM t").collect() result = batches[0].column(0) + # pyarrow.parquet promotes timestamp[s] to timestamp[ms] on write + # (https://github.com/apache/arrow/issues/41382). Compensate so the + # comparison checks DataFusion reads what Arrow actually stored. + if ( + isinstance(arr, pa.Array) + and pa.types.is_timestamp(arr.type) + and arr.type.unit == "s" + ): + arr = arr.cast(pa.timestamp("ms")) + # In DF 43.0.0 we now default to having BinaryView and StringView # so the array that is saved to the parquet is slightly different # than the array read. Convert to values for comparison. diff --git a/python/tests/test_udf.py b/python/tests/test_udf.py index b2540fb57..3a41fa6e1 100644 --- a/python/tests/test_udf.py +++ b/python/tests/test_udf.py @@ -76,6 +76,77 @@ def test_register_udf(ctx, df) -> None: assert result == pa.array([False, False, True]) +def test_udf_lookup(ctx, df) -> None: + is_null = udf( + lambda x: x.is_null(), + [pa.float64()], + pa.bool_(), + volatility="immutable", + name="lookup_is_null", + ) + ctx.register_udf(is_null) + + assert "lookup_is_null" in ctx.udfs() + + looked_up = ctx.udf("lookup_is_null") + df_result = df.select(looked_up(column("b"))) + result = df_result.collect()[0].column(0) + assert result == pa.array([False, False, True]) + + with pytest.raises(KeyError, match="no UDF named"): + ctx.udf("does_not_exist") + + +def test_udf_late_binding_dispatch(ctx, df) -> None: + """Resolve a UDF chosen by configuration string, then invoke it.""" + late_is_null = udf( + lambda x: x.is_null(), + [pa.int64()], + pa.bool_(), + volatility="immutable", + name="late_is_null", + ) + late_is_not_null = udf( + lambda x: pc.invert(x.is_null()), + [pa.int64()], + pa.bool_(), + volatility="immutable", + name="late_is_not_null", + ) + + ctx.register_udf(late_is_null) + ctx.register_udf(late_is_not_null) + + # Pretend this came from a config file / API request — only a string. + runtime_config = {"check_fn": "late_is_not_null"} + + assert runtime_config["check_fn"] in ctx.udfs() + + fn = ctx.udf(runtime_config["check_fn"]) + result = df.select(fn(column("b")).alias("ok")).collect()[0].column(0) + assert result == pa.array([True, True, False]) + + +def test_udaf_lookup_builtin(ctx, df) -> None: + assert "sum" in ctx.udafs() + sum_fn = ctx.udaf("sum") + result = df.aggregate([], [sum_fn(column("a")).alias("total")]).collect() + assert result[0].column(0).to_pylist() == [6] + + with pytest.raises(KeyError, match="no UDAF named"): + ctx.udaf("does_not_exist") + + +def test_udwf_lookup_builtin(ctx, df) -> None: + assert "row_number" in ctx.udwfs() + rn = ctx.udwf("row_number") + result = df.select(column("a"), rn().alias("rn")).collect() + assert result[0].column(1).to_pylist() == [1, 2, 3] + + with pytest.raises(KeyError, match="no UDWF named"): + ctx.udwf("does_not_exist") + + class OverThresholdUDF: def __init__(self, threshold: int = 0) -> None: self.threshold = threshold From fa021fef203c9194747b9ebf1c1f526867d407b1 Mon Sep 17 00:00:00 2001 From: Nick <24689722+ntjohnson1@users.noreply.github.com> Date: Tue, 26 May 2026 11:49:45 -0400 Subject: [PATCH 39/83] Add details on caching to skill (#1521) --- skills/datafusion_python/SKILL.md | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/skills/datafusion_python/SKILL.md b/skills/datafusion_python/SKILL.md index 7b07b430f..98fa2c7aa 100644 --- a/skills/datafusion_python/SKILL.md +++ b/skills/datafusion_python/SKILL.md @@ -253,6 +253,7 @@ polars_df = df.to_polars() # pl.DataFrame py_dict = df.to_pydict() # dict[str, list] py_list = df.to_pylist() # list[dict] count = df.count() # int +df = df.cache() # materialize in memory, return DataFrame ``` ### Date and Timestamp Type Conversion @@ -309,6 +310,29 @@ Async iteration is also supported via `async for batch in df: ...` (or `df.execute_stream()`), which is useful when batches are interleaved with other I/O. +### Caching Intermediate Results + +`df.cache()` materializes a DataFrame as an in-memory table and returns a new +DataFrame backed by it. Reach for it when the same intermediate result feeds +multiple downstream queries — without `cache()`, each branch re-executes the +full upstream plan (re-reading files, recomputing filters/aggregates). + +```python +base = ( + ctx.read_parquet("orders.parquet") + .filter(col("status") == "shipped") + .cache() # materialize once, reuse below +) +by_region = base.aggregate(["region"], [F.sum(col("amount")).alias("total")]) +by_customer = base.aggregate(["customer"], [F.sum(col("amount")).alias("total")]) +``` + +Skip `cache()` for single-use DataFrames — the lazy plan is already optimal. + +The cached table is owned by the DataFrame returned from `cache()` (and any +DataFrames chained from it). To free the memory, drop every reference — let +them go out of scope, or `del base; del by_region; del by_customer`. + ### Writing Results ```python From 081325afe3ac97c6d2ef793a352a26a2634d2738 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 26 May 2026 15:05:53 -0400 Subject: [PATCH 40/83] feat: Python UDFs: per-session inlining toggle and strict refusal setting (#1546) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: per-session Python UDF inlining toggle + sender ctx + strict refusal Adds a per-session toggle that turns inline Python UDF encoding on or off, plus the supporting plumbing to make it usable through pickle.dumps. Codec layer: * PythonLogicalCodec / PythonPhysicalCodec gain a python_udf_inlining bool (default true) and a with_python_udf_inlining(enabled) builder. Each try_encode_udf{,af,wf} short-circuits to inner when the toggle is off; each try_decode_udf{,af,wf} that recognizes a DFPY* magic on a strict codec returns a clean Execution error instead of invoking cloudpickle.loads. The refusal message names the UDF and the wire family so an operator can see at a glance whether to re-encode the bytes or register the UDF on the receiver. Session layer: * PySessionContext::with_python_udf_inlining(enabled) returns a new session whose stacked logical + physical codecs both carry the toggle. The Arc is cloned (cheap), only the codec pair is rebuilt, so registrations and config stay attached. * SessionContext.with_python_udf_inlining(*, enabled) is the Python wrapper. enabled is keyword-only because positional booleans at the call site read as opaque. Sender-side context: * datafusion.ipc gains set_sender_ctx / get_sender_ctx / clear_sender_ctx thread-locals. Expr.__reduce__ now consults get_sender_ctx() to pick the codec for outbound pickles, which is the only path through which a strict session affects pickle.dumps (the protocol calls __reduce__ with no arguments). Without a sender context the default codec is used. Tests: * test_pickle_expr.py picks up TestPythonUdfInliningToggle (covers both directions of the toggle plus the explicit-ctx fast path), TestWorkerCtxLifecycle (set/clear/threading), and TestSenderCtxLifecycle. * New test_pickle_multiprocessing.py + helpers exercise the full driver -> worker round-trip on a multiprocessing.Pool with set_*_ctx installed in the worker initializer. * CI workflow gets a 30-minute timeout-minutes backstop so a hung pickle worker can't block the matrix indefinitely. User-guide docs and the runnable examples land in PR4 of this series. Co-Authored-By: Claude Opus 4.7 (1M context) * update uv lock * docs: clarify Python UDF inlining docstring; drop unresolved :doc: refs Rewrite with_python_udf_inlining docstring for readability and remove references to /user-guide/io/distributing_work, which does not exist yet. Keep security warning inline as a .. warning:: Security block, matching the existing pattern in Expr.to_bytes / from_bytes / __reduce__. The central doc will land in a follow-on PR. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: add doctest examples for sender ctx + UDF inlining toggle Per CLAUDE.md, every Python function needs a docstring example. Adds examples to with_python_udf_inlining, set_sender_ctx, clear_sender_ctx, and get_sender_ctx. Also clarifies that with_python_udf_inlining returns a new SessionContext and leaves the original unchanged, matching the with_logical_extension_codec pattern. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: address review nits for UDF inlining toggle + sender ctx * codec: strict refusal routes through `read_framed_payload` so malformed inline bytes surface their own diagnostic; the "inlining is disabled" message now fires only when the payload would have decoded. * codec: add summary line above `PythonPhysicalCodec::with_python_udf_inlining` cross-link for rustdoc rendering. * expr: hoist `get_sender_ctx` import to module top; note that `__reduce__` also drives `copy.copy` / `copy.deepcopy`. * context: accept `with_python_udf_inlining` positionally or as kwarg (drop `*,`). * tests: replace size-ratio heuristic with semantic check for the `DFPYUDF` family prefix; switch single-batch closure test to `pool.apply`. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: keyword-only inlining flag, skip GIL on prefix mismatch - `SessionContext.with_python_udf_inlining` now keyword-only (`*, enabled`) to match the documented call style and the existing doctests/tests. - `refuse_if_inline` and the three `try_decode_python_*` decoders short- circuit on a `starts_with(family)` check before `Python::attach`, so plans whose UDFs are not Python-defined no longer pay a GIL acquisition per decode call. Semantics preserved: `strip_wire_header` already returns `Ok(None)` when the prefix does not match. - `datafusion.ipc` module docstring wraps the `set_sender_ctx` example in `try`/`finally` and notes that the thread-local holds a strong reference to the installed `SessionContext` until cleared. Co-Authored-By: Claude Opus 4.7 * Add dev dependency * Add testing for CI failure * Additional debugging for mp tests in CI * Set path for workers * more path updates for unit tests * test(pickle): remove multiprocessing CI debug instrumentation Multiprocessing forkserver/spawn hang was diagnosed and fixed: workers could not import `tests._pickle_multiprocessing_helpers` because `pytest --import-mode=importlib` does not add the test parent dir to `sys.path`. The fix (appending the parent dir to `sys.path` so it is inherited by mp workers without shadowing the installed `datafusion` wheel) is retained. This commit drops the diagnostic scaffolding that was added to identify the hang point: - `_diag` + per-import / per-task log writes to /tmp - `snapshot_processes` and the `threading.Timer` that captured worker state mid-hang - `diag_init` Pool initializer - "Dump multiprocessing diagnostic log" CI step Pre-existing infrastructure is kept: per-test `@pytest.mark.timeout(120)` (backed by `pytest-timeout` dev dep) and the job-level `timeout-minutes: 30` backstop on the test matrix. Co-Authored-By: Claude Opus 4.7 * Shorten rust side docstring since it's duplicative of the exposed python docstring * docs: clarify strict-mode refusal message and to_bytes inlining docs Address PR review feedback: - codec.rs: rewrite strict-refusal error to present the two real remediations (sender re-encode by-name + receiver register; or receiver enables inlining, accepting cloudpickle risk) instead of bundling registration with both-side inlining. - expr.py: qualify to_bytes docstring so Python UDF self-contained behavior is conditional on with_python_udf_inlining being enabled. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: clarify with_python_udf_inlining enabled arg is required Reword docstring to drop misleading "(the default)" claim. The `enabled` parameter is keyword-only and required — there is no argument default. Note instead that fresh sessions inline UDFs until the toggle overrides them (a session-level default, not an argument default). Co-Authored-By: Claude Opus 4.7 (1M context) * docs: demonstrate strict-mode refusal in with_python_udf_inlining docstring Replace placeholder isinstance check with a doctest that registers a Python UDF, encodes an expression on the default session, then shows the strict session refusing to decode the inline payload. Exercises the actual behavior the toggle controls. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: convert sender-ctx example to executable doctest Replace the code-block in the ipc module docstring that demonstrated set_sender_ctx with a doctest that actually runs. Worker-init example remains a code-block since it documents a Pool-initializer pattern that does not fit naturally into a doctest. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: use 'thread-local sender context' as adjectival phrase Bare 'thread-local' as a noun reads ambiguously next to the _local.ctx attribute name. Hyphenate as adjective with explicit 'sender context' noun so the referent is unambiguous. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: drop trailing clear_sender_ctx from set_sender_ctx example The trailing cleanup call was test hygiene, not API teaching, and risked implying callers must always pair set with clear. Adjacent clear_sender_ctx and get_sender_ctx doctests are self-contained (they explicitly set or clear before asserting), so removing the cleanup line does not affect doctest outcomes. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .github/workflows/test.yml | 3 + crates/core/src/codec.rs | 157 +++++++++-- crates/core/src/context.rs | 16 ++ pyproject.toml | 1 + python/datafusion/context.py | 58 ++++ python/datafusion/expr.py | 33 ++- python/datafusion/ipc.py | 93 ++++++- .../tests/_pickle_multiprocessing_helpers.py | 89 ++++++ python/tests/test_pickle_expr.py | 254 +++++++++++++++++- python/tests/test_pickle_multiprocessing.py | 145 ++++++++++ uv.lock | 43 ++- 11 files changed, 847 insertions(+), 45 deletions(-) create mode 100644 python/tests/_pickle_multiprocessing_helpers.py create mode 100644 python/tests/test_pickle_multiprocessing.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c597ab308..2cd792ea9 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -29,6 +29,9 @@ env: jobs: test-matrix: runs-on: ubuntu-latest + # Backstop: a hung multiprocessing worker (e.g. during a pickle regression) + # should not block CI longer than this. + timeout-minutes: 30 strategy: fail-fast: false matrix: diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 363ee82b8..b1b9f99dc 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -232,16 +232,39 @@ fn strip_wire_header<'a>( #[derive(Debug)] pub struct PythonLogicalCodec { inner: Arc, + python_udf_inlining: bool, } impl PythonLogicalCodec { pub fn new(inner: Arc) -> Self { - Self { inner } + Self { + inner, + python_udf_inlining: true, + } } pub fn inner(&self) -> &Arc { &self.inner } + + /// Toggle inline encoding of Python UDFs. See + /// `SessionContext.with_python_udf_inlining` (Python) for full + /// behavior and use cases. + /// + /// Security scope: strict mode (`false`) narrows only the codec + /// layer — it stops `Expr::from_bytes` from invoking + /// `cloudpickle.loads` on the inline `DFPY*` payload. It does + /// **not** make `pickle.loads(untrusted_bytes)` safe; treat every + /// `pickle.loads` on untrusted input as unsafe regardless of this + /// setting. + pub fn with_python_udf_inlining(mut self, enabled: bool) -> Self { + self.python_udf_inlining = enabled; + self + } + + pub fn python_udf_inlining(&self) -> bool { + self.python_udf_inlining + } } impl Default for PythonLogicalCodec { @@ -301,48 +324,104 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { - if try_encode_python_scalar_udf(node, buf)? { + if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } self.inner.try_encode_udf(node, buf) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { - if let Some(udf) = try_decode_python_scalar_udf(buf)? { - return Ok(udf); + if self.python_udf_inlining { + if let Some(udf) = try_decode_python_scalar_udf(buf)? { + return Ok(udf); + } + } else { + refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } self.inner.try_decode_udf(name, buf) } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { - if try_encode_python_udaf(node, buf)? { + if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } self.inner.try_encode_udaf(node, buf) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { - if let Some(udaf) = try_decode_python_udaf(buf)? { - return Ok(udaf); + if self.python_udf_inlining { + if let Some(udaf) = try_decode_python_udaf(buf)? { + return Ok(udaf); + } + } else { + refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } self.inner.try_decode_udaf(name, buf) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { - if try_encode_python_udwf(node, buf)? { + if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } self.inner.try_encode_udwf(node, buf) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { - if let Some(udwf) = try_decode_python_udwf(buf)? { - return Ok(udwf); + if self.python_udf_inlining { + if let Some(udwf) = try_decode_python_udwf(buf)? { + return Ok(udwf); + } + } else { + refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } self.inner.try_decode_udwf(name, buf) } } +/// Strict-mode gate: if `buf` is a well-framed inline payload for +/// `family`, return the strict-refusal error; otherwise return +/// `Ok(())` so the caller can delegate to its `inner` codec. +/// +/// Routing through [`read_framed_payload`] (rather than a bare +/// `starts_with` probe) means malformed inline bytes — wrong +/// wire-format version, mismatched Python version, truncated header — +/// surface *their* diagnostic instead of the strict-mode message. +/// The strict message implies sender intent ("inlining is disabled"), +/// so it should fire only when the bytes really would have decoded. +/// +/// Fast path: short-circuit on the family-magic prefix before +/// acquiring the GIL. Plans with many non-Python UDFs would otherwise +/// pay a GIL acquisition per decode call just to confirm "not a +/// Python UDF". `read_framed_payload` itself rejects buffers that +/// don't start with `family`, so this is purely an optimization. +fn refuse_if_inline(buf: &[u8], family: &[u8], kind: &str, name: &str) -> Result<()> { + if !buf.starts_with(family) { + return Ok(()); + } + Python::attach(|py| match read_framed_payload(py, buf, family, kind)? { + Some(_) => Err(refuse_inline_payload(kind, name)), + None => Ok(()), + }) +} + +/// Build the error returned by a strict codec when it receives an +/// inline Python-UDF payload it has been told not to deserialize. +fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusionError { + // `Execution`, not `Plan`: this is a wire-format decode refusal at + // codec time, not a planner-stage failure. Downstream error + // classification keys off the variant — surfacing this as a planner + // error would mis-route it into "fix your SQL" buckets. + datafusion::error::DataFusionError::Execution(format!( + "Refusing to deserialize inline Python {kind} '{name}': Python UDF \ + inlining is disabled on this session. Two remediations: \ + (1) ask the sender to re-encode with inlining disabled so '{name}' \ + travels by name, and register '{name}' on this receiver; or \ + (2) enable inlining on this receiver (accepts the cloudpickle \ + execution risk on inbound payloads). Receivers cannot re-encode \ + bytes they did not produce." + )) +} + /// `PhysicalExtensionCodec` mirror of [`PythonLogicalCodec`] parked /// on the same `SessionContext`. Carries the Python-aware encoding /// hooks for physical-layer types (`ExecutionPlan`, `PhysicalExpr`) @@ -358,16 +437,33 @@ impl LogicalExtensionCodec for PythonLogicalCodec { #[derive(Debug)] pub struct PythonPhysicalCodec { inner: Arc, + python_udf_inlining: bool, } impl PythonPhysicalCodec { pub fn new(inner: Arc) -> Self { - Self { inner } + Self { + inner, + python_udf_inlining: true, + } } pub fn inner(&self) -> &Arc { &self.inner } + + /// Toggle inline encoding of Python UDFs on this physical codec. + /// + /// Mirrors [`PythonLogicalCodec::with_python_udf_inlining`]; see + /// that method for the full security and portability discussion. + pub fn with_python_udf_inlining(mut self, enabled: bool) -> Self { + self.python_udf_inlining = enabled; + self + } + + pub fn python_udf_inlining(&self) -> bool { + self.python_udf_inlining + } } impl Default for PythonPhysicalCodec { @@ -391,15 +487,19 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { - if try_encode_python_scalar_udf(node, buf)? { + if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } self.inner.try_encode_udf(node, buf) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { - if let Some(udf) = try_decode_python_scalar_udf(buf)? { - return Ok(udf); + if self.python_udf_inlining { + if let Some(udf) = try_decode_python_scalar_udf(buf)? { + return Ok(udf); + } + } else { + refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } self.inner.try_decode_udf(name, buf) } @@ -417,29 +517,37 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { - if try_encode_python_udaf(node, buf)? { + if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } self.inner.try_encode_udaf(node, buf) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { - if let Some(udaf) = try_decode_python_udaf(buf)? { - return Ok(udaf); + if self.python_udf_inlining { + if let Some(udaf) = try_decode_python_udaf(buf)? { + return Ok(udaf); + } + } else { + refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } self.inner.try_decode_udaf(name, buf) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { - if try_encode_python_udwf(node, buf)? { + if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } self.inner.try_encode_udwf(node, buf) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { - if let Some(udwf) = try_decode_python_udwf(buf)? { - return Ok(udwf); + if self.python_udf_inlining { + if let Some(udwf) = try_decode_python_udwf(buf)? { + return Ok(udwf); + } + } else { + refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } self.inner.try_decode_udwf(name, buf) } @@ -476,6 +584,9 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec) /// the caller to delegate to its `inner` codec (and eventually the /// `FunctionRegistry`). pub(crate) fn try_decode_python_scalar_udf(buf: &[u8]) -> Result>> { + if !buf.starts_with(PY_SCALAR_UDF_FAMILY) { + return Ok(None); + } Python::attach(|py| -> Result>> { let Some(payload) = read_framed_payload(py, buf, PY_SCALAR_UDF_FAMILY, "scalar UDF")? else { @@ -732,6 +843,9 @@ pub(crate) fn try_encode_python_udwf(node: &WindowUDF, buf: &mut Vec) -> Res } pub(crate) fn try_decode_python_udwf(buf: &[u8]) -> Result>> { + if !buf.starts_with(PY_WINDOW_UDF_FAMILY) { + return Ok(None); + } Python::attach(|py| -> Result>> { let Some(payload) = read_framed_payload(py, buf, PY_WINDOW_UDF_FAMILY, "window UDF")? else { @@ -814,6 +928,9 @@ pub(crate) fn try_encode_python_udaf(node: &AggregateUDF, buf: &mut Vec) -> } pub(crate) fn try_decode_python_udaf(buf: &[u8]) -> Result>> { + if !buf.starts_with(PY_AGG_UDF_FAMILY) { + return Ok(None); + } Python::attach(|py| -> Result>> { let Some(payload) = read_framed_payload(py, buf, PY_AGG_UDF_FAMILY, "aggregate UDF")? else { diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index bca8e30f5..4606246cf 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1453,6 +1453,22 @@ impl PySessionContext { physical_codec, }) } + + pub fn with_python_udf_inlining(&self, enabled: bool) -> Self { + let logical_codec = Arc::new( + PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner())) + .with_python_udf_inlining(enabled), + ); + let physical_codec = Arc::new( + PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) + .with_python_udf_inlining(enabled), + ); + Self { + ctx: Arc::clone(&self.ctx), + logical_codec, + physical_codec, + } + } } impl PySessionContext { diff --git a/pyproject.toml b/pyproject.toml index a02f4608a..418640a49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -203,6 +203,7 @@ dev = [ "pyarrow>=19.0.0", "pygithub==2.5.0", "pytest-asyncio>=0.23.3", + "pytest-timeout>=2.3.1", "pytest>=7.4.4", "pyyaml>=6.0.3", "ruff>=0.15.1", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 9ecbfe311..52bd600c3 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1963,3 +1963,61 @@ def with_physical_extension_codec( new = SessionContext.__new__(SessionContext) new.ctx = new_internal return new + + def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext: + """Control whether Python UDFs are embedded in serialized expressions. + + ``enabled`` is keyword-only and required: callers must pick a + mode explicitly. Fresh sessions inline UDFs (``enabled=True`` + behavior) until this method overrides the toggle. + + With ``enabled=True``, serialized expressions carry the Python + code for any scalar, aggregate, or window UDFs they reference. + The receiver rebuilds the UDFs from those bytes and does not + need to register them first. + + With ``enabled=False``, serialized expressions store only the + UDF names. This has two uses: + + * **Cross-language portability.** The bytes can be decoded by a + non-Python receiver, which must already have UDFs registered + under matching names. + * **Safer deserialization.** :meth:`Expr.from_bytes` will refuse + to rebuild Python UDFs rather than call ``cloudpickle.loads`` + on untrusted input. + + The setting affects :meth:`Expr.to_bytes` and + :meth:`Expr.from_bytes` whenever this session is passed as the + ``ctx`` argument. :func:`pickle.dumps` and :func:`pickle.loads` + do not pass a context, so to apply the setting through pickle, + register this session with + :func:`datafusion.ipc.set_sender_ctx` on the sender and + :func:`datafusion.ipc.set_worker_ctx` on the receiver. + + .. warning:: Security + This setting narrows only :meth:`Expr.from_bytes`. Calling + :func:`pickle.loads` on untrusted bytes remains unsafe + regardless of the toggle. + + Returns a new :class:`SessionContext` with the toggle applied; + the original session is unchanged. + + Examples: + >>> import pyarrow as pa + >>> from datafusion import SessionContext, Expr, col, udf + >>> ctx = SessionContext() + >>> identity = udf(lambda a: a, [pa.int64()], pa.int64(), + ... volatility="immutable", name="identity_demo") + >>> ctx.register_udf(identity) + >>> blob = identity(col("x")).to_bytes(ctx) + >>> strict = SessionContext().with_python_udf_inlining(enabled=False) + >>> try: + ... Expr.from_bytes(blob, strict) + ... except Exception as e: + ... print("Refusing to deserialize" in str(e)) + True + """ + new_internal = self.ctx.with_python_udf_inlining(enabled) + new = SessionContext.__new__(SessionContext) + new.ctx = new_internal + return new diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 7e95bc127..4fdbdc5d4 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -53,6 +53,7 @@ from ._internal import expr as expr_internal from ._internal import functions as functions_internal +from .ipc import get_sender_ctx if TYPE_CHECKING: from collections.abc import Sequence @@ -446,13 +447,18 @@ def to_bytes(self, ctx: SessionContext | None = None) -> bytes: worker process for distributed evaluation. When ``ctx`` is supplied, encoding routes through that session's - installed :class:`LogicalExtensionCodec`. When ``ctx`` is - ``None``, the default codec is used. - - Built-in functions and Python UDFs (scalar, aggregate, window) - travel inside the returned bytes; the worker does not need to - pre-register them. UDFs imported via the FFI capsule protocol - travel by name only and must be registered on the worker. + installed :class:`LogicalExtensionCodec` (so settings like + :meth:`SessionContext.with_python_udf_inlining` take effect). + When ``ctx`` is ``None``, the default codec is used (Python UDF + inlining on, no user-installed extension codec). + + Built-in functions travel inside the returned bytes. Python UDFs + (scalar, aggregate, window) also inline by default, so the worker + does not need to pre-register them; when the encoding session has + :meth:`SessionContext.with_python_udf_inlining` set to ``False``, + Python UDFs travel by name only and must be registered on the + worker. UDFs imported via the FFI capsule protocol always travel + by name only and must be registered on the worker. .. warning:: Security Bytes returned here may embed a cloudpickled Python @@ -522,7 +528,9 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr: Accepts output of :meth:`to_bytes` or :func:`pickle.dumps`. ``ctx`` is the :class:`SessionContext` used to resolve any - function references that travel by name (e.g. FFI UDFs). When + function references that travel by name (e.g. FFI UDFs, or + Python UDFs sent with inlining disabled via + :meth:`SessionContext.with_python_udf_inlining`). When ``ctx`` is ``None`` the worker context installed via :func:`datafusion.ipc.set_worker_ctx` is consulted; if no worker context is installed, the global :class:`SessionContext` is used @@ -586,8 +594,15 @@ def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]: >>> e = col("a") * lit(2) >>> pickle.loads(pickle.dumps(e)).canonical_name() 'a * Int64(2)' + + The encoding side honors a driver-side sender context installed + via :func:`datafusion.ipc.set_sender_ctx` — that is how + :meth:`SessionContext.with_python_udf_inlining` propagates + through ``pickle.dumps``. The sender context is read by + ``__reduce__``, so :func:`copy.copy` and :func:`copy.deepcopy` + — which also go through ``__reduce__`` — pick it up too. """ - return (Expr._reconstruct, (self.to_bytes(),)) + return (Expr._reconstruct, (self.to_bytes(get_sender_ctx()),)) @classmethod def _reconstruct(cls, proto_bytes: bytes) -> Expr: diff --git a/python/datafusion/ipc.py b/python/datafusion/ipc.py index 8dd7fc463..487abd4c3 100644 --- a/python/datafusion/ipc.py +++ b/python/datafusion/ipc.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -"""Worker-side setup for distributing DataFusion expressions. +"""Driver- and worker-side setup for distributing DataFusion expressions. When a :class:`Expr` is shipped to a worker process (e.g. through :func:`multiprocessing.Pool` or a Ray actor), the worker reconstructs the @@ -53,6 +53,36 @@ def init_worker(): payloads are not portable across Python minor versions. See :meth:`datafusion.Expr.to_bytes` for examples of what travels by value vs. by reference. + +On the driver side, call :func:`set_sender_ctx` to control how +:func:`pickle.dumps` encodes expressions — for example, to apply +:meth:`SessionContext.with_python_udf_inlining` to every pickled +expression on this thread: + +>>> import pickle +>>> from datafusion import SessionContext, col, lit +>>> from datafusion.ipc import clear_sender_ctx, set_sender_ctx +>>> driver_ctx = SessionContext().with_python_udf_inlining(enabled=False) +>>> set_sender_ctx(driver_ctx) +>>> try: +... blob = pickle.dumps(col("a") + lit(1)) +... finally: +... clear_sender_ctx() +>>> isinstance(blob, bytes) +True + +Without a sender context the default codec is used (Python UDF +inlining on). The sender context only affects pickle / ``to_bytes`` +encoding; explicit ``expr.to_bytes(ctx)`` calls still use the supplied +``ctx``. + +The thread-local sender context holds a strong reference to the +installed :class:`SessionContext` until :func:`clear_sender_ctx` is +called or the thread exits. Long-running driver threads that install a sender +context once and never clear it will retain that session for the +lifetime of the thread; pair :func:`set_sender_ctx` with +:func:`clear_sender_ctx` (e.g. in a ``try``/``finally``) when the +sender context is only needed for a bounded scope. """ from __future__ import annotations @@ -65,8 +95,11 @@ def init_worker(): __all__ = [ + "clear_sender_ctx", "clear_worker_ctx", + "get_sender_ctx", "get_worker_ctx", + "set_sender_ctx", "set_worker_ctx", ] @@ -125,6 +158,64 @@ def get_worker_ctx() -> SessionContext | None: return getattr(_local, "ctx", None) +def set_sender_ctx(ctx: SessionContext) -> None: + """Install this driver's :class:`SessionContext` for outbound pickles. + + Controls how :func:`pickle.dumps` encodes :class:`Expr` instances on + this thread. The most useful application is propagating a session + configured with + :meth:`SessionContext.with_python_udf_inlining` so the toggle takes + effect through pickle (which otherwise calls + :meth:`Expr.to_bytes` with no context and uses the default codec). + + Idempotent: overwrites any previous value. Stored in a thread-local + slot, so worker threads on the driver may install their own contexts. + Does not affect :meth:`Expr.to_bytes` calls that pass an explicit + ``ctx`` — those continue to use the supplied context. + + Examples: + >>> from datafusion import SessionContext + >>> from datafusion.ipc import set_sender_ctx, get_sender_ctx + >>> driver = SessionContext().with_python_udf_inlining(enabled=False) + >>> set_sender_ctx(driver) + >>> get_sender_ctx() is driver + True + """ + _local.sender_ctx = ctx + + +def clear_sender_ctx() -> None: + """Remove this driver's installed sender :class:`SessionContext`. + + After clearing, pickled expressions fall back to the default codec + (Python UDF inlining on). + + Examples: + >>> from datafusion import SessionContext + >>> from datafusion.ipc import ( + ... set_sender_ctx, clear_sender_ctx, get_sender_ctx, + ... ) + >>> set_sender_ctx(SessionContext()) + >>> clear_sender_ctx() + >>> get_sender_ctx() is None + True + """ + if hasattr(_local, "sender_ctx"): + del _local.sender_ctx + + +def get_sender_ctx() -> SessionContext | None: + """Return this driver's installed sender :class:`SessionContext`, or ``None``. + + Examples: + >>> from datafusion.ipc import get_sender_ctx, clear_sender_ctx + >>> clear_sender_ctx() + >>> get_sender_ctx() is None + True + """ + return getattr(_local, "sender_ctx", None) + + def _resolve_ctx( explicit_ctx: SessionContext | None = None, ) -> SessionContext: diff --git a/python/tests/_pickle_multiprocessing_helpers.py b/python/tests/_pickle_multiprocessing_helpers.py new file mode 100644 index 000000000..4f04967f2 --- /dev/null +++ b/python/tests/_pickle_multiprocessing_helpers.py @@ -0,0 +1,89 @@ +# 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. + +# The leading underscore is load-bearing: pytest with --import-mode=importlib +# (used in CI) assigns synthetic module names to test modules, which breaks +# subprocess imports during multiprocessing. An underscore-prefixed module is +# not collected as a test module, so it imports under its normal __name__ +# inside worker processes. + +from __future__ import annotations + +import pyarrow as pa +from datafusion import SessionContext, udf +from datafusion.ipc import clear_worker_ctx, set_worker_ctx + + +def make_double_udf(): + """Build the canonical UDF used in the multiprocessing tests.""" + return udf( + lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]), + [pa.int64()], + pa.int64(), + volatility="immutable", + name="double", + ) + + +def make_times_seven_udf(): + """Closure-capturing UDF — verifies cloudpickle preserves closed-over state.""" + multiplier = 7 + + def fn(arr): + return pa.array([(v.as_py() or 0) * multiplier for v in arr]) + + return udf( + fn, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="times_seven", + ) + + +def init_worker_empty(): + """Pool initializer: install an empty SessionContext (no UDFs).""" + set_worker_ctx(SessionContext()) + + +def init_worker_clear(): + """Pool initializer: explicitly clear any prior worker context.""" + clear_worker_ctx() + + +def unpickle_and_describe(blob: bytes) -> str: + """Unpickle a proto-bytes blob and return its canonical name.""" + import pickle + + expr = pickle.loads(blob) # noqa: S301 + return expr.canonical_name() + + +def unpickle_and_evaluate(blob: bytes, batch: list[int]) -> list[int]: + """Unpickle an expression and evaluate it against an in-memory batch. + + Returns the result column as a Python list. Used to verify that + cloudpickled UDFs (including closure state) execute correctly in + a fresh worker process. + """ + import pickle + + expr = pickle.loads(blob) # noqa: S301 + ctx = SessionContext() + df = ctx.from_pydict({"a": batch}) + out = df.with_column("result", expr).select("result") + return out.to_pydict()["result"] diff --git a/python/tests/test_pickle_expr.py b/python/tests/test_pickle_expr.py index eb0441c49..588caa21a 100644 --- a/python/tests/test_pickle_expr.py +++ b/python/tests/test_pickle_expr.py @@ -21,27 +21,36 @@ with the pickled expression and do not need worker-side pre-registration. The worker context (:mod:`datafusion.ipc`) is only consulted for UDFs imported via the FFI capsule protocol. + +Cross-process tests live in ``test_pickle_multiprocessing.py``. """ from __future__ import annotations import pickle +import threading import pyarrow as pa import pytest from datafusion import Expr, SessionContext, col, lit, udf from datafusion.ipc import ( + clear_sender_ctx, clear_worker_ctx, + get_sender_ctx, + get_worker_ctx, + set_sender_ctx, set_worker_ctx, ) @pytest.fixture(autouse=True) def _reset_worker_ctx(): - """Ensure every test starts with no worker context installed.""" + """Ensure every test starts with no worker or sender context installed.""" clear_worker_ctx() + clear_sender_ctx() yield clear_worker_ctx() + clear_sender_ctx() def _double_udf(): @@ -124,6 +133,8 @@ def fn(arr): e = u(col("a")) blob = pickle.dumps(e) decoded = pickle.loads(blob) # noqa: S301 + # Round-trip names match; functional verification of captured state + # happens in test_pickle_multiprocessing via an actual UDF call. assert decoded.canonical_name() == e.canonical_name() def test_multi_arg_udf_round_trip(self): @@ -311,3 +322,244 @@ def test_cross_version_error_message(self): Exception, match="not portable across Python minor versions" ): Expr.from_bytes(bytes(tampered)) + + +class TestPythonUdfInliningToggle: + """`SessionContext.with_python_udf_inlining(enabled=False)` opts out of + inline Python UDF encoding for both encode and decode paths.""" + + def _build_double_udf(self): + return udf( + lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]), + [pa.int64()], + pa.int64(), + volatility="immutable", + name="double", + ) + + def test_strict_encoder_omits_inline_payload(self): + """Strict mode emits the by-name wire form: no `DFPYUDF` magic + in the blob, no cloudpickled callable. Semantic check is + sharper than a size-ratio heuristic — a renamed UDF or a + smaller-than-expected closure would still flip the magic + bytes, but might not move the size by 4x. + """ + ctx_inline = SessionContext() + ctx_strict = ctx_inline.with_python_udf_inlining(enabled=False) + u = self._build_double_udf() + e = u(col("a")) + + blob_inline = e.to_bytes(ctx_inline) + blob_strict = e.to_bytes(ctx_strict) + + # `DFPYUDF` is the scalar Python-UDF family prefix; see + # `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs. + assert b"DFPYUDF" in blob_inline + assert b"DFPYUDF" not in blob_strict + + def test_toggle_off_then_on_restores_inline_encoding(self): + """`with_python_udf_inlining` is per-call clone semantics: + flipping off and then on must produce a context that emits the + same inline form as a fresh default context, byte-for-byte. + + Guards against a regression where the off→on transition leaves + the codec in a sticky strict state (e.g. by mutating shared + codec state instead of cloning). + """ + u = self._build_double_udf() + e = u(col("a")) + + baseline = SessionContext() + toggled = ( + SessionContext() + .with_python_udf_inlining(enabled=False) + .with_python_udf_inlining(enabled=True) + ) + + blob_baseline = e.to_bytes(baseline) + blob_toggled = e.to_bytes(toggled) + + assert blob_baseline == blob_toggled + + # Sanity check the decoded form against a fresh ctx — the + # toggled-back blob should be self-contained inline, not a + # strict by-name payload that needs registry resolution. + decoded = Expr.from_bytes(blob_toggled, ctx=SessionContext()) + assert "double" in decoded.canonical_name() + + def test_strict_roundtrip_via_registry(self): + """When both sender and receiver disable inlining, the UDF + travels by name only and the receiver resolves it from its + registered functions.""" + strict_sender = SessionContext().with_python_udf_inlining(enabled=False) + u = self._build_double_udf() + blob = u(col("a")).to_bytes(strict_sender) + + receiver = SessionContext().with_python_udf_inlining(enabled=False) + receiver.register_udf(u) + restored = Expr.from_bytes(blob, ctx=receiver) + assert "double" in restored.canonical_name() + + def test_strict_decoder_refuses_inline_payload(self): + """An inline-encoded blob fed to a strict receiver raises with a + clear error rather than silently invoking cloudpickle.loads. + + The receiver is intentionally *not* given a matching + registration: the codec refusal must trip before the registry + is ever consulted, so registering the UDF here would only mask + a regression that moved the check after registry lookup. + """ + sender = SessionContext() + u = self._build_double_udf() + blob = u(col("a")).to_bytes(sender) + + strict_receiver = SessionContext().with_python_udf_inlining(enabled=False) + # `RuntimeError` (not bare `Exception`): the codec refusal is + # surfaced through `parse_expr` → `PyRuntimeError`. Tightening + # the assertion catches a regression that swallows the refusal + # as a different error type. + with pytest.raises(RuntimeError, match="inlining is disabled"): + Expr.from_bytes(blob, ctx=strict_receiver) + + def test_sender_ctx_propagates_through_pickle(self): + """`set_sender_ctx` makes `pickle.dumps` use a strict codec. + + Without a sender context, pickle defaults to the inline codec + and the blob contains the `DFPYUDF` family prefix. With a + strict sender context installed, the callable encodes by name + and the prefix is absent. + """ + u = self._build_double_udf() + e = u(col("a")) + + blob_default = pickle.dumps(e) + + strict_sender = SessionContext().with_python_udf_inlining(enabled=False) + set_sender_ctx(strict_sender) + try: + blob_strict = pickle.dumps(e) + finally: + clear_sender_ctx() + + assert b"DFPYUDF" in blob_default + assert b"DFPYUDF" not in blob_strict + + def test_sender_ctx_strict_roundtrip_via_pickle(self): + """End-to-end pickle round-trip with strict mode on both sides. + + Driver installs a strict sender context. Worker installs a + matching strict context with the UDF registered. The UDF + travels by name through `pickle.dumps` / `pickle.loads`. + """ + u = self._build_double_udf() + e = u(col("a")) + + strict_sender = SessionContext().with_python_udf_inlining(enabled=False) + set_sender_ctx(strict_sender) + try: + blob = pickle.dumps(e) + finally: + clear_sender_ctx() + + worker = SessionContext().with_python_udf_inlining(enabled=False) + worker.register_udf(u) + set_worker_ctx(worker) + try: + decoded = pickle.loads(blob) # noqa: S301 + finally: + clear_worker_ctx() + + assert "double" in decoded.canonical_name() + + def test_sender_ctx_strict_pickle_accepted_by_inline_worker_with_registry(self): + """A strict-encoded blob still decodes fine on an inline worker + because the wire format is the same default-codec by-name form. + Sanity check: cross-config works as long as the receiver can + resolve the name.""" + u = self._build_double_udf() + e = u(col("a")) + + strict_sender = SessionContext().with_python_udf_inlining(enabled=False) + set_sender_ctx(strict_sender) + try: + blob = pickle.dumps(e) + finally: + clear_sender_ctx() + + worker = SessionContext() + worker.register_udf(u) + set_worker_ctx(worker) + try: + decoded = pickle.loads(blob) # noqa: S301 + finally: + clear_worker_ctx() + + assert "double" in decoded.canonical_name() + + +class TestWorkerCtxLifecycle: + def test_set_and_clear(self): + assert get_worker_ctx() is None + ctx = SessionContext() + set_worker_ctx(ctx) + assert get_worker_ctx() is ctx + clear_worker_ctx() + assert get_worker_ctx() is None + + def test_clear_when_unset_is_noop(self): + clear_worker_ctx() # no error + assert get_worker_ctx() is None + + def test_thread_local_isolation(self): + main_ctx = SessionContext() + set_worker_ctx(main_ctx) + + seen_in_thread: list = [] + + def worker(): + seen_in_thread.append(get_worker_ctx()) + set_worker_ctx(SessionContext()) + seen_in_thread.append(get_worker_ctx()) + + t = threading.Thread(target=worker) + t.start() + t.join() + + # Thread saw no ctx initially (thread-local), then its own. + assert seen_in_thread[0] is None + assert seen_in_thread[1] is not main_ctx + # Main thread's ctx is unchanged by the thread's actions. + assert get_worker_ctx() is main_ctx + + +class TestSenderCtxLifecycle: + def test_set_and_clear(self): + assert get_sender_ctx() is None + ctx = SessionContext() + set_sender_ctx(ctx) + assert get_sender_ctx() is ctx + clear_sender_ctx() + assert get_sender_ctx() is None + + def test_clear_when_unset_is_noop(self): + clear_sender_ctx() # no error + assert get_sender_ctx() is None + + def test_thread_local_isolation(self): + main_ctx = SessionContext() + set_sender_ctx(main_ctx) + + seen_in_thread: list = [] + + def worker(): + seen_in_thread.append(get_sender_ctx()) + set_sender_ctx(SessionContext()) + seen_in_thread.append(get_sender_ctx()) + + t = threading.Thread(target=worker) + t.start() + t.join() + + assert seen_in_thread[0] is None + assert seen_in_thread[1] is not main_ctx + assert get_sender_ctx() is main_ctx diff --git a/python/tests/test_pickle_multiprocessing.py b/python/tests/test_pickle_multiprocessing.py new file mode 100644 index 000000000..fcce49d97 --- /dev/null +++ b/python/tests/test_pickle_multiprocessing.py @@ -0,0 +1,145 @@ +# 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. + +"""Cross-process pickle tests for :class:`Expr`. + +Workers run with each :mod:`multiprocessing` start method (``fork``, +``forkserver``, ``spawn``). Python UDFs (scalar, aggregate, window) travel +with the pickled expression and need no worker-side pre-registration. +Worker-side helpers live in ``_pickle_multiprocessing_helpers`` — the +underscore prefix avoids pytest collection so the module imports under +its real name in worker subprocesses. +""" + +from __future__ import annotations + +import functools +import multiprocessing as mp +import pickle +import sys +from pathlib import Path + +import pytest +from datafusion import col, lit + +from . import _pickle_multiprocessing_helpers as helpers + +# `pytest --import-mode=importlib` (used in CI) does not put the test parent +# directory on `sys.path`; pytest loads `tests` via its own importlib hook. +# multiprocessing forkserver / spawn workers receive only the parent's +# `sys.path` snapshot, not pytest's hook, so they fail to import +# `tests._pickle_multiprocessing_helpers` with `ModuleNotFoundError: No module +# named 'tests'`. Append (not prepend) the parent directory of the `tests` +# package so workers can resolve it the standard way, *without* shadowing the +# installed `datafusion` wheel — the source tree's `python/datafusion/` has +# no `_internal` extension module (that lives in the wheel under +# site-packages), so prepending would break `from datafusion._internal +# import ...`. Fork start method is unaffected (inherits the already-imported +# module object). +_TESTS_PARENT = str(Path(__file__).resolve().parent.parent) +if _TESTS_PARENT not in sys.path: + sys.path.append(_TESTS_PARENT) + + +@functools.cache +def _multiprocessing_available() -> tuple[bool, str]: + """Return (available, reason). Some sandboxed environments deny semaphore + creation; without semaphores, ``multiprocessing.Pool`` cannot start. + + Cached so the probe Pool only spawns once per session, and only when a + test in this module is actually about to run — collection-only runs + (e.g. ``pytest --collect-only`` on the full suite) skip the probe. + """ + try: + ctx = mp.get_context("spawn") + with ctx.Pool(processes=1) as pool: + pool.map(int, [0]) + except (PermissionError, OSError) as exc: + return False, f"multiprocessing.Pool unavailable: {exc}" + return True, "" + + +@pytest.fixture(autouse=True) +def _skip_if_multiprocessing_unavailable(): + available, reason = _multiprocessing_available() + if not available: + pytest.skip(reason) + + +START_METHODS = [ + pytest.param( + "fork", + marks=pytest.mark.skipif( + sys.platform == "darwin", + reason="fork start method is unsafe with PyArrow/tokio on macOS", + ), + ), + "forkserver", + "spawn", +] + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.timeout(120) +def test_builtin_pickle_via_pool(start_method): + """Built-in expressions round-trip in every start method.""" + expr = col("a") + lit(1) + blob = pickle.dumps(expr) + + ctx = mp.get_context(start_method) + with ctx.Pool(processes=2) as pool: + results = pool.map(helpers.unpickle_and_describe, [blob, blob, blob]) + + assert all(r == expr.canonical_name() for r in results) + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.timeout(120) +def test_udf_pickle_self_contained(start_method): + """Scalar UDF travels inside the proto blob — no worker pre-registration. + + Workers start with no UDF registered. The Rust-side ``PythonUDFCodec`` + reconstructs the UDF from bytes embedded in the pickle blob. + """ + udf_obj = helpers.make_double_udf() + expr = udf_obj(col("a")) + blob = pickle.dumps(expr) + + ctx = mp.get_context(start_method) + with ctx.Pool(processes=2) as pool: + results = pool.starmap( + helpers.unpickle_and_evaluate, + [(blob, [1, 2, 3]), (blob, [10, 20, 30])], + ) + + assert results[0] == [2, 4, 6] + assert results[1] == [20, 40, 60] + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.timeout(120) +def test_closure_capturing_udf_via_pool(start_method): + """Cloudpickle preserves closure state across the codec boundary.""" + udf_obj = helpers.make_times_seven_udf() + expr = udf_obj(col("a")) + blob = pickle.dumps(expr) + + ctx = mp.get_context(start_method) + with ctx.Pool(processes=2) as pool: + result = pool.apply(helpers.unpickle_and_evaluate, (blob, [1, 2, 3])) + + assert result == [7, 14, 21] diff --git a/uv.lock b/uv.lock index 3fd3eec4b..26ab8b20e 100644 --- a/uv.lock +++ b/uv.lock @@ -343,6 +343,7 @@ dev = [ { name = "pygithub" }, { name = "pytest" }, { name = "pytest-asyncio" }, + { name = "pytest-timeout" }, { name = "pyyaml" }, { name = "ruff" }, { name = "toml" }, @@ -380,6 +381,7 @@ dev = [ { name = "pygithub", specifier = "==2.5.0" }, { name = "pytest", specifier = ">=7.4.4" }, { name = "pytest-asyncio", specifier = ">=0.23.3" }, + { name = "pytest-timeout", specifier = ">=2.3.1" }, { name = "pyyaml", specifier = ">=6.0.3" }, { name = "ruff", specifier = ">=0.15.1" }, { name = "toml", specifier = ">=0.10.2" }, @@ -628,25 +630,26 @@ wheels = [ [[package]] name = "maturin" -version = "1.8.1" +version = "1.13.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "tomli", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9a/08/ccb0f917722a35ab0d758be9bb5edaf645c3a3d6170061f10d396ecd273f/maturin-1.8.1.tar.gz", hash = "sha256:49cd964aabf59f8b0a6969f9860d2cdf194ac331529caae14c884f5659568857", size = 197397, upload-time = "2024-12-30T14:03:48.109Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/1c/612d23d33ec21b9ae7ece7b3f0dd5f9dfd57b4009e9d2938165869ebd6ae/maturin-1.13.3.tar.gz", hash = "sha256:771e1e9e71a278e56db01552e0d1acfd1464259f9575b6e72842f893cd299079", size = 357934, upload-time = "2026-05-11T07:43:39.027Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4c/00/f34077315f34db8ad2ccf6bfe11b864ca27baab3a1320634da8e3cf89a48/maturin-1.8.1-py3-none-linux_armv6l.whl", hash = "sha256:7e590a23d9076b8a994f2e67bc63dc9a2d1c9a41b1e7b45ac354ba8275254e89", size = 7568415, upload-time = "2024-12-30T14:03:07.939Z" }, - { url = "https://files.pythonhosted.org/packages/5c/07/9219976135ce0cb32d2fa6ea5c6d0ad709013d9a17967312e149b98153a6/maturin-1.8.1-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:8d8251a95682c83ea60988c804b620c181911cd824aa107b4a49ac5333c92968", size = 14527816, upload-time = "2024-12-30T14:03:13.851Z" }, - { url = "https://files.pythonhosted.org/packages/e6/04/fa009a00903acdd1785d58322193140bfe358595347c39f315112dabdf9e/maturin-1.8.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:b9fc1a4354cac5e32c190410208039812ea88c4a36bd2b6499268ec49ef5de00", size = 7580446, upload-time = "2024-12-30T14:03:17.64Z" }, - { url = "https://files.pythonhosted.org/packages/9b/d4/414b2aab9bbfe88182b734d3aa1b4fef7d7701e50f6be48500378b8c8721/maturin-1.8.1-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:621e171c6b39f95f1d0df69a118416034fbd59c0f89dcaea8c2ea62019deecba", size = 7650535, upload-time = "2024-12-30T14:03:21.115Z" }, - { url = "https://files.pythonhosted.org/packages/f0/64/879418a8a0196013ec1fb19eada0781c04a30e8d6d9227e80f91275a4f5b/maturin-1.8.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:98f638739a5132962347871b85c91f525c9246ef4d99796ae98a2031e3df029f", size = 8006702, upload-time = "2024-12-30T14:03:24.318Z" }, - { url = "https://files.pythonhosted.org/packages/39/c2/605829324f8371294f70303aca130682df75318958efed246873d3d604ab/maturin-1.8.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:f9f5c47521924b6e515cbc652a042fe5f17f8747445be9d931048e5d8ddb50a4", size = 7368164, upload-time = "2024-12-30T14:03:26.582Z" }, - { url = "https://files.pythonhosted.org/packages/be/6c/30e136d397bb146b94b628c0ef7f17708281611b97849e2cf37847025ac7/maturin-1.8.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:0f4407c7353c31bfbb8cdeb82bc2170e474cbfb97b5ba27568f440c9d6c1fdd4", size = 7450889, upload-time = "2024-12-30T14:03:28.893Z" }, - { url = "https://files.pythonhosted.org/packages/1b/50/e1f5023512696d4e56096f702e2f68d6d9a30afe0a4eec82b0e27b8eb4e4/maturin-1.8.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:ec49cd70cad3c389946c6e2bc0bd50772a7fcb463040dd800720345897eec9bf", size = 9585819, upload-time = "2024-12-30T14:03:31.125Z" }, - { url = "https://files.pythonhosted.org/packages/b7/80/b24b5248d89d2e5982553900237a337ea098ca9297b8369ca2aa95549e0f/maturin-1.8.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c08767d794de8f8a11c5c8b1b47a4ff9fb6ae2d2d97679e27030f2f509c8c2a0", size = 10920801, upload-time = "2024-12-30T14:03:35.127Z" }, - { url = "https://files.pythonhosted.org/packages/6e/f4/8ede7a662fabf93456b44390a5ad22630e25fb5ddaecf787251071b2e143/maturin-1.8.1-py3-none-win32.whl", hash = "sha256:d678407713f3e10df33c5b3d7a343ec0551eb7f14d8ad9ba6febeb96f4e4c75c", size = 6873556, upload-time = "2024-12-30T14:03:37.913Z" }, - { url = "https://files.pythonhosted.org/packages/9c/22/757f093ed0e319e9648155b8c9d716765442bea5bc98ebc58ad4ad5b0524/maturin-1.8.1-py3-none-win_amd64.whl", hash = "sha256:a526f90fe0e5cb59ffb81f4ff547ddc42e823bbdeae4a31012c0893ca6dcaf46", size = 7823153, upload-time = "2024-12-30T14:03:40.33Z" }, - { url = "https://files.pythonhosted.org/packages/a4/f5/051413e04f6da25069db5e76759ecdb8cd2a8ab4a94045b5a3bf548c66fa/maturin-1.8.1-py3-none-win_arm64.whl", hash = "sha256:e95f077fd2ddd2f048182880eed458c308571a534be3eb2add4d3dac55bf57f4", size = 6552131, upload-time = "2024-12-30T14:03:45.203Z" }, + { url = "https://files.pythonhosted.org/packages/71/66/18c2aaac0b2a5dea9f1db5984ce83b905ad205cfc7c02d0091e707c0c2e7/maturin-1.13.3-py3-none-linux_armv6l.whl", hash = "sha256:3cc13929ca82aefa4adbf0f2c35419369796213c6fb0eb24e914945f50ef5d8c", size = 10190971, upload-time = "2026-05-11T07:43:10.431Z" }, + { url = "https://files.pythonhosted.org/packages/bc/71/26a988d092e4fd6a9523d46d44400a46cad7cdf3fd206ce702240c748aee/maturin-1.13.3-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:53b08bd075649ce96513ad9abf241a43cb685ed6e9e7790f8dbc2d66e95d8323", size = 19716714, upload-time = "2026-05-11T07:43:36.911Z" }, + { url = "https://files.pythonhosted.org/packages/82/5c/f3fd0e184255d9fc7e272c62af3dfa84c617b2577ef83af9ce615f5279cc/maturin-1.13.3-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:4cd478e6e4c56251e48ed079b8efd55b30bc5c09cf695a1bdafaeb582ee735a0", size = 10194726, upload-time = "2026-05-11T07:43:07.05Z" }, + { url = "https://files.pythonhosted.org/packages/a9/e1/f4edb69fb647b77c4769a9bfd4d6fb62961e653d164bc277ecdffac3ab61/maturin-1.13.3-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:a2675e25f313034ae6f57388cf14818f87d8961c4a96795287f3e155f59beb11", size = 10172781, upload-time = "2026-05-11T07:43:40.796Z" }, + { url = "https://files.pythonhosted.org/packages/c7/7d/a1be934690cdcc3c6609769ceaad322ab7501c2ee5bafcac1b14d609e403/maturin-1.13.3-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:4667ef609ab446c1b5e0bfe4f9fb99699ab6d8548433f8d1a684256e0b67217f", size = 10682670, upload-time = "2026-05-11T07:43:13.132Z" }, + { url = "https://files.pythonhosted.org/packages/18/f5/372ae19b72ce8f6e37e5864ae4dc5b252ee9fce0619ccc3aa366aa3a7f97/maturin-1.13.3-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:3db93337ed97e60ffc878aa8b493cd7ae44d3a5e1a37256db3a4491f57565018", size = 10060363, upload-time = "2026-05-11T07:43:21.107Z" }, + { url = "https://files.pythonhosted.org/packages/cb/5b/c68340cca09368af0df80965dfabed4234205a492a93da00793c7b9aae20/maturin-1.13.3-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:1cc0a110b224ca90406b668a3e3c1f5a515062e59e26292f6dbaf5fd4909c6f3", size = 10017551, upload-time = "2026-05-11T07:43:33.916Z" }, + { url = "https://files.pythonhosted.org/packages/28/1e/f90fb2b000bad9e6d850cd5afb88b2f1e2a279cfb4de02ea40078484690e/maturin-1.13.3-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:c00ea6428dea17bf616fe93770837634454b28c2de1a876e42ef8036c616079a", size = 13301712, upload-time = "2026-05-11T07:43:26.492Z" }, + { url = "https://files.pythonhosted.org/packages/be/58/1670f68a8f04ccd7b90df11047bd9a046585310e84e1967cc9849cd1c5a3/maturin-1.13.3-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:49fd6ab08da28098ccf37afca24cdba72376ba9c1eedf9dd25ff82ed771961ff", size = 10946765, upload-time = "2026-05-11T07:43:16.135Z" }, + { url = "https://files.pythonhosted.org/packages/4b/ac/00c955c2ef134817b1a7bdaa76b0309e9c5291eb17d9ff88069eecd08bc2/maturin-1.13.3-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:b6741d7bf4af97da937528fd1e523c6ab54f53d9a21870fa735d6e67fd88e273", size = 10388661, upload-time = "2026-05-11T07:43:18.727Z" }, + { url = "https://files.pythonhosted.org/packages/97/c6/cbf8a51dde19c19aeba0d9b075095a2effb9b31fd312b1aae3ac79f8aea2/maturin-1.13.3-py3-none-win32.whl", hash = "sha256:0ef257e692cc756c87af5bea95ddfe7d3ac49d3376a7a87f728d63f06e7b6f8b", size = 8901838, upload-time = "2026-05-11T07:43:23.76Z" }, + { url = "https://files.pythonhosted.org/packages/a1/ff/c6a50a59dc8313097d43ac5f4d74df6a500c8cb62b0dc9e054f53e203a48/maturin-1.13.3-py3-none-win_amd64.whl", hash = "sha256:def4a435ea9d2ee93b18ba579dc8c9cf898889a66f312cd379b5e374ec3e3ad6", size = 10340801, upload-time = "2026-05-11T07:43:29.239Z" }, + { url = "https://files.pythonhosted.org/packages/6c/93/e32e79333f0902ba292b996f504f5f06be59587f7d02ab8d5ed1e3066445/maturin-1.13.3-py3-none-win_arm64.whl", hash = "sha256:2389fe92d017cea9d94e521fa0175314a4c52f79a1057b901fbc9f8686ef7d0b", size = 9706562, upload-time = "2026-05-11T07:43:31.743Z" }, ] [[package]] @@ -1238,6 +1241,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/67/17/3493c5624e48fd97156ebaec380dcaafee9506d7e2c46218ceebbb57d7de/pytest_asyncio-0.25.3-py3-none-any.whl", hash = "sha256:9e89518e0f9bd08928f97a3482fdc4e244df17529460bc038291ccaf8f85c7c3", size = 19467, upload-time = "2025-01-28T18:37:56.798Z" }, ] +[[package]] +name = "pytest-timeout" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/82/4c9ecabab13363e72d880f2fb504c5f750433b2b6f16e99f4ec21ada284c/pytest_timeout-2.4.0.tar.gz", hash = "sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a", size = 17973, upload-time = "2025-05-05T19:44:34.99Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/b6/3127540ecdf1464a00e5a01ee60a1b09175f6913f0644ac748494d9c4b21/pytest_timeout-2.4.0-py3-none-any.whl", hash = "sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2", size = 14382, upload-time = "2025-05-05T19:44:33.502Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" From 56b1ceaae2a5023c6b3238999dc55fcee781aa24 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 27 May 2026 13:22:19 -0400 Subject: [PATCH 41/83] Bump DataFusion to 1321d60 (54.0.0) (#1562) Update the pinned DataFusion git rev to 1321d60cc37ee487d1e7ce7f501357c3236b2542, which is DataFusion 54.0.0. Bump the workspace dependency requirements from 53 to 54 so the [patch.crates-io] git overrides actually bind (cargo only applies a patch when its version satisfies the dependency requirement), and refresh Cargo.lock accordingly. Adapt to the 54 API: - Remove the DatasetExec::apply_expressions override; apply_expressions is no longer a member of the ExecutionPlan trait. - factorial now errors on negative input, so take abs() before applying factorial in the parametrized expr test and update the expected values. Co-authored-by: Claude --- Cargo.lock | 541 +++++++++++++++----------------- Cargo.toml | 36 +-- crates/core/src/dataset_exec.rs | 10 +- python/tests/test_expr.py | 7 +- 4 files changed, 278 insertions(+), 316 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0c4b77582..6a1ef2447 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -78,9 +78,9 @@ dependencies = [ [[package]] name = "arc-swap" -version = "1.9.0" +version = "1.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a07d1f37ff60921c83bdfc7407723bdefe89b44b98a9b772f225c8f9d67141a6" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" dependencies = [ "rustversion", ] @@ -167,7 +167,7 @@ dependencies = [ "flate2", "indexmap", "liblzma", - "rand 0.9.2", + "rand 0.9.4", "serde", "serde_json", "snap", @@ -361,9 +361,9 @@ dependencies = [ [[package]] name = "async-compression" -version = "0.4.41" +version = "0.4.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0f9ee0f6e02ffd7ad5816e9464499fba7b3effd01123b515c41d1697c43dad1" +checksum = "e79b3f8a79cccc2898f31920fc69f304859b3bd567490f75ebf51ae1c792a9ac" dependencies = [ "compression-codecs", "compression-core", @@ -416,9 +416,9 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "autocfg" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "base64" @@ -441,9 +441,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.11.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" [[package]] name = "blake2" @@ -456,16 +456,16 @@ dependencies = [ [[package]] name = "blake3" -version = "1.8.3" +version = "1.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2468ef7d57b3fb7e16b576e8377cdbde2320c60e1491e961d11da40fc4f02a2d" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" dependencies = [ "arrayref", "arrayvec", "cc", "cfg-if", "constant_time_eq", - "cpufeatures 0.2.17", + "cpufeatures", ] [[package]] @@ -509,9 +509,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.20.2" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "byteorder" @@ -536,9 +536,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.58" +version = "1.2.62" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1e928d4b69e3077709075a938a05ffbedfa53a84c8f766efbf8220bb1ff60e1" +checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98" dependencies = [ "find-msvc-tools", "jobserver", @@ -565,8 +565,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", - "rand_core 0.10.0", + "cpufeatures", + "rand_core 0.10.1", ] [[package]] @@ -612,9 +612,9 @@ dependencies = [ [[package]] name = "compression-codecs" -version = "0.4.37" +version = "0.4.38" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb7b51a7d9c967fc26773061ba86150f19c50c0d65c887cb1fbe295fd16619b7" +checksum = "ce2548391e9c1929c21bf6aa2680af86fe4c1b33e6cea9ac1cfeec0bd11218cf" dependencies = [ "bzip2", "compression-core", @@ -627,9 +627,9 @@ dependencies = [ [[package]] name = "compression-core" -version = "0.4.31" +version = "0.4.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75984efb6ed102a0d42db99afb6c1948f0380d1d91808d5529916e6c08b49d8d" +checksum = "cc14f565cf027a105f7a44ccf9e5b424348421a1d8952a8fc9d499d313107789" [[package]] name = "const-oid" @@ -679,15 +679,6 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "cpufeatures" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" -dependencies = [ - "libc", -] - [[package]] name = "cpufeatures" version = "0.3.0" @@ -745,9 +736,9 @@ dependencies = [ [[package]] name = "crypto-common" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" dependencies = [ "hybrid-array", ] @@ -785,9 +776,9 @@ dependencies = [ [[package]] name = "dashmap" -version = "6.1.0" +version = "6.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" dependencies = [ "cfg-if", "crossbeam-utils", @@ -799,8 +790,8 @@ dependencies = [ [[package]] name = "datafusion" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-schema", @@ -852,8 +843,8 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "async-trait", @@ -876,8 +867,8 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "async-trait", @@ -898,8 +889,8 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-ipc", @@ -923,8 +914,8 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "futures", "log", @@ -933,8 +924,8 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "async-compression", @@ -959,7 +950,7 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.2", + "rand 0.9.4", "tokio", "tokio-util", "url", @@ -968,8 +959,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-ipc", @@ -991,8 +982,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-avro", @@ -1009,8 +1000,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "async-trait", @@ -1031,8 +1022,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "async-trait", @@ -1053,8 +1044,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "async-trait", @@ -1083,13 +1074,13 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" [[package]] name = "datafusion-execution" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-buffer", @@ -1102,15 +1093,15 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.2", + "rand 0.9.4", "tempfile", "url", ] [[package]] name = "datafusion-expr" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-schema", @@ -1131,8 +1122,8 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1142,8 +1133,8 @@ dependencies = [ [[package]] name = "datafusion-ffi" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-schema", @@ -1196,8 +1187,8 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-buffer", @@ -1219,7 +1210,7 @@ dependencies = [ "md-5 0.11.0", "memchr", "num-traits", - "rand 0.9.2", + "rand 0.9.4", "regex", "sha2", "uuid", @@ -1227,8 +1218,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1247,8 +1238,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1258,8 +1249,8 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-ord", @@ -1282,22 +1273,23 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "async-trait", "datafusion-catalog", "datafusion-common", "datafusion-expr", + "datafusion-physical-expr", "datafusion-physical-plan", "parking_lot", ] [[package]] name = "datafusion-functions-window" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1312,8 +1304,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1321,8 +1313,8 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "datafusion-doc", "quote", @@ -1331,8 +1323,8 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "chrono", @@ -1350,8 +1342,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1371,8 +1363,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1385,8 +1377,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "chrono", @@ -1401,8 +1393,8 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1419,8 +1411,8 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "arrow-data", @@ -1451,8 +1443,8 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "chrono", @@ -1477,8 +1469,8 @@ dependencies = [ [[package]] name = "datafusion-proto-common" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1487,8 +1479,8 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "datafusion-common", @@ -1546,8 +1538,8 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "async-trait", "datafusion-common", @@ -1559,8 +1551,8 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "arrow", "bigdecimal", @@ -1577,8 +1569,8 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "53.1.0" -source = "git+https://github.com/apache/datafusion?rev=47655fd6c9ef060d73497987e6ccb98e57196508#47655fd6c9ef060d73497987e6ccb98e57196508" +version = "54.0.0" +source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" dependencies = [ "async-recursion", "async-trait", @@ -1613,14 +1605,14 @@ checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ "block-buffer 0.12.0", "const-oid", - "crypto-common 0.2.1", + "crypto-common 0.2.2", ] [[package]] name = "displaydoc" -version = "0.2.5" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", @@ -1635,9 +1627,9 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "either" -version = "1.15.0" +version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" [[package]] name = "equivalent" @@ -1657,9 +1649,9 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.3.0" +version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" [[package]] name = "find-msvc-tools" @@ -1855,7 +1847,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", - "rand_core 0.10.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -1868,9 +1860,9 @@ checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "h2" -version = "0.4.13" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "171fefbc92fe4a4de27e0698d6a5b392d6a0e333506bc49133760b3bcf948733" dependencies = [ "atomic-waker", "bytes", @@ -1948,9 +1940,9 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "http" -version = "1.4.0" +version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" +checksum = "8be7462df143984c4598a256ef469b251d7d7f9e271135073e78fc535414f3d0" dependencies = [ "bytes", "itoa", @@ -2002,9 +1994,9 @@ dependencies = [ [[package]] name = "hyper" -version = "1.8.1" +version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11" +checksum = "eb92f162bf56536459fc83c79b974bb12837acfed43d6bc370a7916d0ae15ecc" dependencies = [ "atomic-waker", "bytes", @@ -2016,7 +2008,6 @@ dependencies = [ "httparse", "itoa", "pin-project-lite", - "pin-utils", "smallvec", "tokio", "want", @@ -2024,16 +2015,15 @@ dependencies = [ [[package]] name = "hyper-rustls" -version = "0.27.7" +version = "0.27.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ "http", "hyper", "hyper-util", "rustls", "rustls-native-certs", - "rustls-pki-types", "tokio", "tokio-rustls", "tower-service", @@ -2088,12 +2078,13 @@ dependencies = [ [[package]] name = "icu_collections" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" dependencies = [ "displaydoc", "potential_utf", + "utf8_iter", "yoke", "zerofrom", "zerovec", @@ -2101,9 +2092,9 @@ dependencies = [ [[package]] name = "icu_locale_core" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" dependencies = [ "displaydoc", "litemap", @@ -2114,9 +2105,9 @@ dependencies = [ [[package]] name = "icu_normalizer" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" dependencies = [ "icu_collections", "icu_normalizer_data", @@ -2128,15 +2119,15 @@ dependencies = [ [[package]] name = "icu_normalizer_data" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" [[package]] name = "icu_properties" -version = "2.1.2" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" dependencies = [ "icu_collections", "icu_locale_core", @@ -2148,15 +2139,15 @@ dependencies = [ [[package]] name = "icu_properties_data" -version = "2.1.2" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" [[package]] name = "icu_provider" -version = "2.1.1" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" dependencies = [ "displaydoc", "icu_locale_core", @@ -2186,9 +2177,9 @@ dependencies = [ [[package]] name = "idna_adapter" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" dependencies = [ "icu_normalizer", "icu_properties", @@ -2218,16 +2209,6 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" -[[package]] -name = "iri-string" -version = "0.7.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8e7418f59cc01c88316161279a7f665217ae316b388e58a0d10e29f54f1e5eb" -dependencies = [ - "memchr", - "serde", -] - [[package]] name = "itertools" version = "0.14.0" @@ -2255,10 +2236,12 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.91" +version = "0.3.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b49715b7073f385ba4bc528e5747d02e66cb39c6146efb66b781f131f0fb399c" +checksum = "142bc4740e452c1e57ade0cbc129f139c9093e354346f0872ef985f4f5cf5f11" dependencies = [ + "cfg-if", + "futures-util", "once_cell", "wasm-bindgen", ] @@ -2328,9 +2311,9 @@ dependencies = [ [[package]] name = "libbz2-rs-sys" -version = "0.2.2" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c4a545a15244c7d945065b5d392b2d2d7f21526fba56ce51467b06ed445e8f7" +checksum = "34b357333733e8260735ba5894eb928c02ecc69c78715f01a8019e7fa7f2db4c" [[package]] name = "libc" @@ -2359,9 +2342,9 @@ dependencies = [ [[package]] name = "liblzma-sys" -version = "0.4.5" +version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f2db66f3268487b5033077f266da6777d057949b8f93c8ad82e441df25e6186" +checksum = "1a60851d15cd8c5346eca4ab8babff585be2ae4bc8097c067291d3ffe2add3b6" dependencies = [ "cc", "libc", @@ -2376,12 +2359,11 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.44" +version = "0.1.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" dependencies = [ "cc", - "libc", ] [[package]] @@ -2392,9 +2374,9 @@ checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" [[package]] name = "litemap" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "lock_api" @@ -2407,9 +2389,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.29" +version = "0.4.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +checksum = "616ec5685824bcc94416c6d4a7a446eea774a31efd7062c8480ba6fd06d7a6e5" [[package]] name = "lru-slab" @@ -2419,9 +2401,9 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.13.0" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db9a0d582c2874f68138a16ce1867e0ffde6c0bb0a0df85e1f36d04146db488a" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" dependencies = [ "twox-hash", ] @@ -2448,15 +2430,15 @@ dependencies = [ [[package]] name = "memchr" -version = "2.8.0" +version = "2.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +checksum = "6b947ae49db0d222b1dbc6b113ce7248a3fc3a6ca21b696717bfc000ba4484d8" [[package]] name = "mimalloc" -version = "0.1.48" +version = "0.1.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1ee66a4b64c74f4ef288bcbb9192ad9c3feaad75193129ac8509af543894fd8" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" dependencies = [ "libmimalloc-sys", ] @@ -2559,7 +2541,7 @@ dependencies = [ "parking_lot", "percent-encoding", "quick-xml", - "rand 0.10.0", + "rand 0.10.1", "reqwest", "ring", "rustls-pki-types", @@ -2736,18 +2718,18 @@ dependencies = [ [[package]] name = "pin-project" -version = "1.1.12" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbf0d9e68100b3a7989b4901972f265cd542e560a3a8a724e1e20322f4d06ce9" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.12" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a990e22f43e84855daf260dded30524ef4a9021cc7541c26540500a50b624389" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", @@ -2760,17 +2742,11 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" -[[package]] -name = "pin-utils" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" - [[package]] name = "pkg-config" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" [[package]] name = "portable-atomic" @@ -2780,9 +2756,9 @@ checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" [[package]] name = "potential_utf" -version = "0.1.4" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" dependencies = [ "zerovec", ] @@ -2886,9 +2862,9 @@ dependencies = [ [[package]] name = "psm" -version = "0.1.30" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3852766467df634d74f0b2d7819bf8dc483a0eb2e3b0f50f756f9cfe8b0d18d8" +checksum = "645dbe486e346d9b5de3ef16ede18c26e6c70ad97418f4874b8b1889d6e761ea" dependencies = [ "ar_archive_writer", "cc", @@ -2896,9 +2872,9 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.28.2" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf85e27e86080aafd5a22eae58a162e133a589551542b3e5cee4beb27e54f8e1" +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" dependencies = [ "libc", "once_cell", @@ -2924,18 +2900,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.28.2" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bf94ee265674bf76c09fa430b0e99c26e319c945d96ca0d5a8215f31bf81cf7" +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.28.2" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "491aa5fc66d8059dd44a75f4580a2962c1862a1c2945359db36f6c2818b748dc" +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" dependencies = [ "libc", "pyo3-build-config", @@ -2954,9 +2930,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.28.2" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f5d671734e9d7a43449f8480f8b38115df67bef8d21f76837fa75ee7aaa5e52e" +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -2966,9 +2942,9 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.28.2" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22faaa1ce6c430a1f71658760497291065e6450d7b5dc2bcf254d49f66ee700a" +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" dependencies = [ "heck", "proc-macro2", @@ -2979,9 +2955,9 @@ dependencies = [ [[package]] name = "quick-xml" -version = "0.39.2" +version = "0.39.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "958f21e8e7ceb5a1aa7fa87fab28e7c75976e0bfe7e23ff069e0a260f894067d" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" dependencies = [ "memchr", "serde", @@ -3016,7 +2992,7 @@ dependencies = [ "bytes", "getrandom 0.3.4", "lru-slab", - "rand 0.9.2", + "rand 0.9.4", "ring", "rustc-hash", "rustls", @@ -3076,9 +3052,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.2" +version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.5", @@ -3086,13 +3062,13 @@ dependencies = [ [[package]] name = "rand" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" dependencies = [ "chacha20", "getrandom 0.4.2", - "rand_core 0.10.0", + "rand_core 0.10.1", ] [[package]] @@ -3135,9 +3111,9 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" [[package]] name = "recursive" @@ -3265,9 +3241,9 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.1" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" [[package]] name = "rustc_version" @@ -3293,9 +3269,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.37" +version = "0.23.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" +checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" dependencies = [ "once_cell", "ring", @@ -3319,9 +3295,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.0" +version = "1.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" dependencies = [ "web-time", "zeroize", @@ -3329,9 +3305,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.10" +version = "0.103.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df33b2b81ac578cabaf06b89b0631153a3f416b0a886e8a7a1707fb51abbd1ef" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" dependencies = [ "ring", "rustls-pki-types", @@ -3480,9 +3456,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.149" +version = "1.0.150" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ "indexmap", "itoa", @@ -3536,7 +3512,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures", "digest 0.11.3", ] @@ -3566,9 +3542,9 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" [[package]] name = "siphasher" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" [[package]] name = "slab" @@ -3663,15 +3639,15 @@ checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "stacker" -version = "0.1.23" +version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d74a23609d509411d10e2176dc2a4346e3b4aea2e7b1869f19fdedbc71c013" +checksum = "640c8cdd92b6b12f5bcb1803ca3bbf5ab96e5e6b6b96b9ab77dabe9e880b3190" dependencies = [ "cc", "cfg-if", "libc", "psm", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3822,9 +3798,9 @@ dependencies = [ [[package]] name = "tinystr" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" dependencies = [ "displaydoc", "zerovec", @@ -3917,9 +3893,9 @@ dependencies = [ [[package]] name = "toml_edit" -version = "0.25.11+spec-1.1.0" +version = "0.25.12+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" +checksum = "d2153edc6955a6c354fad8f5efd38b6a8769bdccf9fe50f8e1329f81b0baa5d7" dependencies = [ "indexmap", "toml_datetime", @@ -3953,20 +3929,20 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.8" +version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "bitflags", "bytes", "futures-util", "http", "http-body", - "iri-string", "pin-project-lite", "tower", "tower-layer", "tower-service", + "url", ] [[package]] @@ -4133,9 +4109,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.23.0" +version = "1.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9" +checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" dependencies = [ "getrandom 0.4.2", "js-sys", @@ -4175,11 +4151,11 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] name = "wasip2" -version = "1.0.2+wasi-0.2.9" +version = "1.0.3+wasi-0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" dependencies = [ - "wit-bindgen", + "wit-bindgen 0.57.1", ] [[package]] @@ -4188,14 +4164,14 @@ version = "0.4.0+wasi-0.3.0-rc-2026-01-06" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" dependencies = [ - "wit-bindgen", + "wit-bindgen 0.51.0", ] [[package]] name = "wasm-bindgen" -version = "0.2.114" +version = "0.2.122" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6532f9a5c1ece3798cb1c2cfdba640b9b3ba884f5db45973a6f442510a87d38e" +checksum = "3ed04576f974d2b2fba0f38c51dbc5518011e38c36bf1143164be765528fd409" dependencies = [ "cfg-if", "once_cell", @@ -4206,23 +4182,19 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.64" +version = "0.4.72" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9c5522b3a28661442748e09d40924dfb9ca614b21c00d3fd135720e48b67db8" +checksum = "9473dbd2991ae90b6291c3c32c30c6187ac49aa32f9905d1cce280ec1e110b0f" dependencies = [ - "cfg-if", - "futures-util", "js-sys", - "once_cell", "wasm-bindgen", - "web-sys", ] [[package]] name = "wasm-bindgen-macro" -version = "0.2.114" +version = "0.2.122" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18a2d50fcf105fb33bb15f00e7a77b772945a2ee45dcf454961fd843e74c18e6" +checksum = "916151b09da36bd82f6615cbf3a419e2f0ba23a03c6160e8e92eb6bd4aa1dec6" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4230,9 +4202,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.114" +version = "0.2.122" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03ce4caeaac547cdf713d280eda22a730824dd11e6b8c3ca9e42247b25c631e3" +checksum = "299047362ccbfce148b67ab7e73349f77748e00c8296f9542adfad2ad82c5c5e" dependencies = [ "bumpalo", "proc-macro2", @@ -4243,9 +4215,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.114" +version = "0.2.122" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75a326b8c223ee17883a4251907455a2431acc2791c98c26279376490c378c16" +checksum = "9a929b2c61f11ba3e9bc35b50c1f25cb38e0e892c0c231ae2b8cf78d5dad4437" dependencies = [ "unicode-ident", ] @@ -4299,9 +4271,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.91" +version = "0.3.99" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "854ba17bb104abfb26ba36da9729addc7ce7f06f5c0f90f3c391f8461cca21f9" +checksum = "6d621441cfc37b84979402712047321980c178f299193a3589d05b99e8763436" dependencies = [ "js-sys", "wasm-bindgen", @@ -4394,15 +4366,6 @@ dependencies = [ "windows-targets 0.52.6", ] -[[package]] -name = "windows-sys" -version = "0.59.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets 0.52.6", -] - [[package]] name = "windows-sys" version = "0.60.2" @@ -4552,9 +4515,9 @@ checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" [[package]] name = "winnow" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ee1708bef14716a11bae175f579062d4554d95be2c6829f518df847b7b3fdd0" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" dependencies = [ "memchr", ] @@ -4568,6 +4531,12 @@ dependencies = [ "wit-bindgen-rust-macro", ] +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + [[package]] name = "wit-bindgen-core" version = "0.51.0" @@ -4649,15 +4618,15 @@ dependencies = [ [[package]] name = "writeable" -version = "0.6.2" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" [[package]] name = "yoke" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" +checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" dependencies = [ "stable_deref_trait", "yoke-derive", @@ -4666,9 +4635,9 @@ dependencies = [ [[package]] name = "yoke-derive" -version = "0.8.1" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", @@ -4678,18 +4647,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.47" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efbb2a062be311f2ba113ce66f697a4dc589f85e78a4aea276200804cea0ed87" +checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.47" +version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e8bc7269b54418e7aeeef514aa68f8690b8c0489a06b0136e5f57c4c5ccab89" +checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" dependencies = [ "proc-macro2", "quote", @@ -4698,18 +4667,18 @@ dependencies = [ [[package]] name = "zerofrom" -version = "0.1.6" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" dependencies = [ "zerofrom-derive", ] [[package]] name = "zerofrom-derive" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", @@ -4725,9 +4694,9 @@ checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" [[package]] name = "zerotrie" -version = "0.2.3" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a59c17a5562d507e4b54960e8569ebee33bee890c70aa3fe7b97e85a9fd7851" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" dependencies = [ "displaydoc", "yoke", @@ -4736,9 +4705,9 @@ dependencies = [ [[package]] name = "zerovec" -version = "0.11.5" +version = "0.11.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c28719294829477f525be0186d13efa9a3c602f7ec202ca9e353d310fb9a002" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" dependencies = [ "yoke", "zerofrom", @@ -4747,9 +4716,9 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.11.2" +version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index 13d7040a2..e72c22368 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,15 +40,15 @@ arrow = { version = "58" } arrow-array = { version = "58" } arrow-schema = { version = "58" } arrow-select = { version = "58" } -datafusion = { version = "53" } -datafusion-substrait = { version = "53" } -datafusion-proto = { version = "53" } -datafusion-ffi = { version = "53" } -datafusion-catalog = { version = "53", default-features = false } -datafusion-common = { version = "53", default-features = false } -datafusion-functions-aggregate = { version = "53" } -datafusion-functions-window = { version = "53" } -datafusion-expr = { version = "53" } +datafusion = { version = "54" } +datafusion-substrait = { version = "54" } +datafusion-proto = { version = "54" } +datafusion-ffi = { version = "54" } +datafusion-catalog = { version = "54", default-features = false } +datafusion-common = { version = "54", default-features = false } +datafusion-functions-aggregate = { version = "54" } +datafusion-functions-window = { version = "54" } +datafusion-expr = { version = "54" } prost = "0.14.3" serde_json = "1" uuid = { version = "1.23" } @@ -71,12 +71,12 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] -datafusion = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-common = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } -datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "47655fd6c9ef060d73497987e6ccb98e57196508" } +datafusion = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-common = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } +datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } diff --git a/crates/core/src/dataset_exec.rs b/crates/core/src/dataset_exec.rs index 771119a0f..32c030b00 100644 --- a/crates/core/src/dataset_exec.rs +++ b/crates/core/src/dataset_exec.rs @@ -21,12 +21,11 @@ use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::error::{ArrowError, Result as ArrowResult}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; -use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::{DataFusionError as InnerDataFusionError, Result as DFResult}; use datafusion::execution::context::TaskContext; use datafusion::logical_expr::Expr; use datafusion::logical_expr::utils::conjunction; -use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; +use datafusion::physical_expr::{EquivalenceProperties, LexOrdering}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ @@ -234,13 +233,6 @@ impl ExecutionPlan for DatasetExec { Ok(Arc::new(self.projected_statistics.clone())) } - fn apply_expressions( - &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> DFResult, - ) -> DFResult { - Ok(TreeNodeRecursion::Continue) - } - fn properties(&self) -> &Arc { &self.plan_properties } diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py index 485d69624..05f91acca 100644 --- a/python/tests/test_expr.py +++ b/python/tests/test_expr.py @@ -550,9 +550,10 @@ def test_alias_with_metadata(df): id="atanh", ), pytest.param( - # large numbers cause an integer overflow so divid to make smaller - (col("b") / lit(4)).factorial(), - pa.array([1, 3628800, 1, None], type=pa.int64()), + # large numbers cause an integer overflow so divide to make smaller; + # factorial of a negative number is undefined, so take abs first + (col("b").abs() / lit(4)).factorial(), + pa.array([5040, 3628800, 1, None], type=pa.int64()), id="factorial", ), pytest.param( From 23f9179ad08189637f88218a8ea77a1222262abc Mon Sep 17 00:00:00 2001 From: BharatDeva Date: Thu, 28 May 2026 08:34:43 -0500 Subject: [PATCH 42/83] docs: document null-handling function arguments (#1527) Co-authored-by: BharatDeva <278575558+BharatDeva@users.noreply.github.com> --- python/datafusion/functions.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 28c10e005..383853007 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -910,6 +910,9 @@ def chr(arg: Expr) -> Expr: def coalesce(*args: Expr) -> Expr: """Returns the value of the first expr in ``args`` which is not NULL. + Args: + *args: Expressions to evaluate in order. + Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [None, 1], "b": [2, 3]}) @@ -1091,6 +1094,10 @@ def greatest(*args: Expr) -> Expr: def ifnull(x: Expr, y: Expr) -> Expr: """Returns ``x`` if ``x`` is not NULL. Otherwise returns ``y``. + Args: + x: Expression to return when it is not NULL. + y: Fallback expression to return when ``x`` is NULL. + See Also: This is an alias for :py:func:`nvl`. """ @@ -1325,6 +1332,10 @@ def md5(arg: Expr) -> Expr: def nanvl(x: Expr, y: Expr) -> Expr: """Returns ``x`` if ``x`` is not ``NaN``. Otherwise returns ``y``. + Args: + x: Expression to return when it is not NaN. + y: Fallback expression to return when ``x`` is NaN. + Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [np.nan, 1.0], "b": [0.0, 0.0]}) @@ -1341,6 +1352,10 @@ def nanvl(x: Expr, y: Expr) -> Expr: def nvl(x: Expr, y: Expr) -> Expr: """Returns ``x`` if ``x`` is not ``NULL``. Otherwise returns ``y``. + Args: + x: Expression to return when it is not NULL. + y: Fallback expression to return when ``x`` is NULL. + Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [None, 1], "b": [0, 0]}) @@ -1358,6 +1373,11 @@ def nvl(x: Expr, y: Expr) -> Expr: def nvl2(x: Expr, y: Expr, z: Expr) -> Expr: """Returns ``y`` if ``x`` is not NULL. Otherwise returns ``z``. + Args: + x: Expression to check for NULL. + y: Expression to return when ``x`` is not NULL. + z: Expression to return when ``x`` is NULL. + Examples: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict({"a": [None, 1], "b": [10, 20], "c": [30, 40]}) From baec559b0a7c85934338d6da80ffbe538004f4d4 Mon Sep 17 00:00:00 2001 From: BharatDeva Date: Thu, 28 May 2026 14:28:47 -0500 Subject: [PATCH 43/83] fix: type scalar UDF returns as Arrow arrays (#1528) Co-authored-by: BharatDeva <278575558+BharatDeva@users.noreply.github.com> --- python/datafusion/user_defined.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index a8ee7756e..ced707a96 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: from _typeshed import CapsuleType as _PyCapsule - _R = TypeVar("_R", bound=pa.DataType) + _R = TypeVar("_R", bound=pa.Array) from collections.abc import Callable, Sequence @@ -137,7 +137,7 @@ def __init__( name: str, func: Callable[..., _R], input_fields: list[pa.Field], - return_field: _R, + return_field: pa.Field, volatility: Volatility | str, ) -> None: """Instantiate a scalar user-defined function (UDF). @@ -311,7 +311,7 @@ def _function( def _decorator( input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field, - return_field: _R, + return_field: pa.DataType | pa.Field, volatility: Volatility | str, name: str | None = None, ) -> Callable: From 0a9ca68ade8bbe06aaf67928bb8edc65670baa24 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 28 May 2026 17:09:33 -0400 Subject: [PATCH 44/83] docs: user guide + runnable examples for distributing expressions (#1547) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: user guide page + runnable examples for distributing expressions Wraps up the Expr-pickle work with the user-facing material: * docs/source/user-guide/io/distributing_work.rst — new user guide page covering the multiprocessing, Ray, and datafusion-distributed patterns. Includes the Security section that is the canonical home for the cloudpickle / pickle.loads threat model. * docs/source/user-guide/io/index.rst — toctree entry. * examples/multiprocessing_pickle_expr.py — runnable example: a Pool.map of a closure-capturing UDF across processes, with worker context registration in the initializer. * examples/ray_pickle_expr.py — Ray actor analogue. * examples/datafusion-ffi-example/python/tests/_test_pickle_strict_ffi.py — exercises the strict-mode refusal end to end against an FFI capsule scalar UDF (kept under the FFI example crate because the test needs that crate's compiled artifacts). * examples/README.md — index entries for the new files. Also tightens three docstrings that previously duplicated the security warning so they point at the canonical Security section instead: * PythonLogicalCodec::with_python_udf_inlining (rustdoc): one-line summary plus a relative pointer to distributing_work.rst and the upstream Python pickle module security warning. * SessionContext.with_python_udf_inlining: one-sentence summary plus :doc: link to the user guide. * datafusion.ipc module docstring: cross-reference to the user guide for the full pattern. The crate-level codec.rs module rustdoc also updates "pure-Python scalar UDFs" to "scalar / aggregate / window UDFs" now that all three are covered. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: document Python-version and import portability caveats for inline UDFs Reviewer feedback on the Expr-pickle PRs (#1544) asked that the cloudpickle portability caveats be discoverable on the user-facing page, not only in docstrings. The distributing_work.rst page is the designated canonical home for the distribution story, so add them here: * New 'Portability requirements for inline Python UDFs' subsection covering the matching-Python-minor-version requirement and the by-value vs by-reference import-capture rule (imported modules must be importable on the worker). * Qualify the 'fully portable' Python-UDF bullet to point at the new requirements. * Cross-reference the new subsection from the closure-capture note. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: restore version-byte and cloudpickle-cache rustdoc wording Two codec.rs docstrings were reworded in PR4 in ways that dropped information: * try_encode_python_scalar_udf: restore the `DFPYUDF` family prefix + version byte description of the payload framing (PR4 had collapsed it to `DFPYUDF1` prefix, dropping the version-byte mention). * cloudpickle cached-handle comment: restore "The encode/decode helpers above" wording. * docs: fix reversed tuple order in multiprocessing example docstring The 'Worker layout' docstring described tasks as `(expr, label)` but the code builds and unpacks them as `(label, expr)`. Correct the doc to match. * Respond to first batch of reviewer comments * docs: relocate and restructure distributing-work guide Move the page from user-guide/io/ to the top level of user-guide/ — distributing work is a runtime/operational concern, not a file-format topic, and the shorter "Distributing work" title fits the sidebar cleanly. Restructure the body to lead with the practical worker-setup pattern instead of the four-slot SessionContext taxonomy. The taxonomy survives at the bottom as a reference subsection; the worker-init example and portability rules now reach the reader before they need it. Also addresses reviewer NIT: wrap the `if __name__ == "__main__":` guidance in a `.. note::` admonition and link to the Python multiprocessing docs. Add a header paragraph to each runnable example pointing to the user-guide page so a reader who jumps straight to the example gets the surrounding context. Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/codec.rs | 19 +- docs/source/index.rst | 1 + docs/source/user-guide/distributing-work.rst | 368 ++++++++++++++++++ examples/README.md | 5 + .../python/tests/_test_pickle_strict_ffi.py | 127 ++++++ examples/multiprocessing_pickle_expr.py | 172 ++++++++ examples/ray_pickle_expr.py | 86 ++++ 7 files changed, 771 insertions(+), 7 deletions(-) create mode 100644 docs/source/user-guide/distributing-work.rst create mode 100644 examples/datafusion-ffi-example/python/tests/_test_pickle_strict_ffi.py create mode 100644 examples/multiprocessing_pickle_expr.py create mode 100644 examples/ray_pickle_expr.py diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index b1b9f99dc..a6ea6671c 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -19,11 +19,11 @@ //! //! Datafusion-python plans can carry references to Python-defined //! objects that the upstream protobuf codecs do not know how to -//! serialize: pure-Python scalar UDFs, Python query-planning -//! extensions, and so on. Their state lives inside `Py` -//! callables and closures rather than being recoverable from a name -//! in the receiver's function registry. To ship a plan across a -//! process boundary (pickle, `multiprocessing`, Ray actor, +//! serialize: pure-Python scalar / aggregate / window UDFs, Python +//! query-planning extensions, and so on. Their state lives inside +//! `Py` callables and closures rather than being recoverable +//! from a name in the receiver's function registry. To ship a plan +//! across a process boundary (pickle, `multiprocessing`, Ray actor, //! `datafusion-distributed`, etc.) those payloads have to be encoded //! into the proto wire format itself. //! @@ -256,7 +256,12 @@ impl PythonLogicalCodec { /// `cloudpickle.loads` on the inline `DFPY*` payload. It does /// **not** make `pickle.loads(untrusted_bytes)` safe; treat every /// `pickle.loads` on untrusted input as unsafe regardless of this - /// setting. + /// setting. See `docs/source/user-guide/io/distributing_work.rst` + /// (Security section) for the full threat model, and Python's + /// [pickle module security warning][1] for why `pickle.loads` is + /// unsafe in general. + /// + /// [1]: https://docs.python.org/3/library/pickle.html#module-pickle pub fn with_python_udf_inlining(mut self, enabled: bool) -> Self { self.python_udf_inlining = enabled; self @@ -433,7 +438,7 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio /// encoding on this layer too — otherwise a plan with a Python UDF /// would round-trip at the logical level but break at the physical /// level. Both layers reuse the shared payload framing -/// ([`PY_SCALAR_UDF_FAMILY`]) so the wire format is identical. +/// ([`PY_SCALAR_UDF_FAMILY`] et al.) so the wire format is identical. #[derive(Debug)] pub struct PythonPhysicalCodec { inner: Arc, diff --git a/docs/source/index.rst b/docs/source/index.rst index 0007cc41a..7edb69807 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -76,6 +76,7 @@ Example user-guide/common-operations/index user-guide/io/index user-guide/configuration + user-guide/distributing-work user-guide/sql user-guide/upgrade-guides user-guide/ai-coding-assistants diff --git a/docs/source/user-guide/distributing-work.rst b/docs/source/user-guide/distributing-work.rst new file mode 100644 index 000000000..03b5ca0b9 --- /dev/null +++ b/docs/source/user-guide/distributing-work.rst @@ -0,0 +1,368 @@ +.. 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. + +Distributing work +================= + +DataFusion supports splitting work across processes by shipping +serialized expressions to workers: the driver builds an +:py:class:`~datafusion.Expr`, each worker evaluates it against its +own slice of data. This pattern suits embarrassingly-parallel +workloads where the driver decides partitioning up front. + +Query-level distribution — where the runtime partitions a single +logical or physical plan across worker nodes — is in progress +upstream via `datafusion-distributed +`_ and `Apache +Ballista `_. Both +have short sections at the end of this page; integration details +will land as those projects become usable from datafusion-python. + +Expression-level distribution +----------------------------- + +DataFusion expressions support distribution directly: pass one to a +worker process and Python's standard +`pickle `_ machinery +serializes it transparently — the same machinery +:py:meth:`multiprocessing.pool.Pool.map`, Ray's ``@ray.remote``, and +similar libraries already use to ship function arguments. Python UDFs +— scalar, aggregate, and window — travel inside the serialized +expression; the receiver does not need to pre-register them. + +Basic worker-pool example +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Define a worker function that takes the expression plus a batch and +returns the evaluated result: + +.. code-block:: python + + import pyarrow as pa + from datafusion import SessionContext + + + def evaluate(expr, batch): + # `expr` arrived here via the pool's automatic pickling — + # no manual serialization needed in user code. + ctx = SessionContext() + df = ctx.from_pydict({"a": batch}) + return df.with_column("result", expr).select("result").to_pydict()["result"] + +Then build the expression in the driver and fan it out: + +.. code-block:: python + + import multiprocessing as mp + from datafusion import col, udf + + double = udf( + lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]), + [pa.int64()], pa.int64(), volatility="immutable", name="double", + ) + expr = double(col("a")) + + mp_ctx = mp.get_context("forkserver") + with mp_ctx.Pool(processes=4) as pool: + results = pool.starmap( + evaluate, + [(expr, [1, 2, 3]), (expr, [10, 20, 30])], + ) + print(results) # [[2, 4, 6], [20, 40, 60]] + +.. note:: + + When saved to a ``.py`` file and executed with the ``spawn`` or + ``forkserver`` start method, wrap the driver block in + ``if __name__ == "__main__":`` so worker processes can re-import + the module without re-running it. This is a standard Python + :py:mod:`multiprocessing` requirement, not DataFusion-specific — + see `Safe importing of main module + `_ + in the Python docs. + + +What travels with the expression +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +* **Built-in functions** (``abs``, ``length``, arithmetic, comparisons, + etc.) — fully portable. Worker needs nothing pre-registered. +* **Python UDFs** — travel inline (subject to the two portability + requirements below). The callable, its signature, and any state + captured in closures travel inside the serialized expression and are + reconstructed on the worker automatically. Applies equally to: + + * **scalar UDFs** (:py:func:`datafusion.udf`) + * **aggregate UDFs** (:py:func:`datafusion.udaf`) + * **window UDFs** (:py:func:`datafusion.udwf`) +* **UDFs imported via the FFI capsule protocol** — travel **by name + only**. The worker must already have a matching registration on its + :py:class:`SessionContext`. Without that registration, evaluation + raises an error. + +Portability requirements for inline Python UDFs +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Inline Python UDFs ride on `cloudpickle +`_, which imposes two +requirements on the worker environment: + +* **Matching Python minor version.** cloudpickle serializes Python + bytecode, which is not stable across minor versions. A UDF pickled + on 3.12 cannot be reconstructed on 3.11 or 3.13. The wire format + stamps the sender's ``(major, minor)``; mismatches raise a clear + error naming both versions. Align the Python version on driver and + workers. +* **Imported modules must be importable on the worker.** cloudpickle + captures the callable *by value* (bytecode and closure cells travel + whole), but names resolved through ``import`` are captured *by + reference* — module path only. A UDF doing + ``from mylib import transform`` requires ``mylib`` installed on the + worker. Same applies to bound methods of imported classes. + Self-contained UDFs (no imports beyond what the worker already has, + e.g. ``pyarrow``) avoid this entirely. + +Registering shared UDFs on workers +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When an expression references an FFI capsule UDF (or any UDF the +worker must resolve from its registered functions), set up the +worker's :py:class:`SessionContext` once per process and install it +as the *worker context*: + +.. code-block:: python + + from datafusion import SessionContext + from datafusion.ipc import set_worker_ctx + + + def init_worker(): + ctx = SessionContext() + ctx.register_udaf(my_ffi_aggregate) + set_worker_ctx(ctx) + + + with mp.get_context("forkserver").Pool( + processes=4, initializer=init_worker + ) as pool: + ... + +Inside a worker, expressions arriving from the driver resolve their +by-name references against the installed worker context. If no worker +context is installed, the global :py:class:`SessionContext` is used — +fine for expressions that only reference built-ins and Python UDFs, +but FFI-capsule-backed registrations must be installed on the global +context to resolve. + +Python 3.14 default change +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Python 3.14 changed the Linux default start method for +:py:mod:`multiprocessing` from ``fork`` to ``forkserver`` (macOS has +defaulted to ``spawn`` since Python 3.8; Windows has always used +``spawn``). With ``fork``, any state set in the parent was visible in +workers via copy-on-write; with ``forkserver`` and ``spawn`` it is +not. The :py:func:`~datafusion.ipc.set_worker_ctx` pattern works on +every start method — prefer it over relying on inherited state. + +Practical considerations +~~~~~~~~~~~~~~~~~~~~~~~~ + +* **Serialized size scales with what travels inline.** A serialized + expression of just built-ins is small (tens of bytes). An + expression carrying a Python UDF is hundreds of bytes (the callable + and its signature). When the same UDF is shipped many times, + registering an equivalent FFI-capsule UDF on each worker via + :py:func:`~datafusion.ipc.set_worker_ctx` and referring to it by + name cuts the per-trip overhead. +* **Closure capture.** When a Python UDF closes over surrounding + state — local variables, module-level objects, file paths — that + state is captured at serialization time. Surprises are possible if + the captured state is large, mutable, or not portable to the + worker's environment. See `Portability requirements for inline + Python UDFs`_ for the Python-version and imported-module rules. + +Disabling Python UDF inlining +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For a stricter wire format, call +:py:meth:`SessionContext.with_python_udf_inlining(enabled=False) +` on the session +producing or consuming the bytes. With inlining disabled, Python +UDFs travel by name only — the same way FFI-capsule UDFs do — and +the receiver must have a matching registration. + +Two use cases: + +* **Cross-language portability.** A non-Python decoder cannot + reconstruct a cloudpickled payload. Senders aimed at Java, C++, + or another Rust binary disable inlining and rely on the receiver + having compatible UDF registrations. +* **Untrusted-source decode.** With inlining disabled, + :py:meth:`Expr.from_bytes` never calls ``cloudpickle.loads`` on + the incoming bytes — an inline payload from a misbehaving sender + raises a clear error instead of executing arbitrary Python code. + +Mismatched configurations raise a descriptive error: an inline blob +fed to a strict receiver fails fast rather than silently dropping +into ``cloudpickle.loads``. + +To make the toggle apply through :py:func:`pickle.dumps` (which +calls :py:meth:`Expr.to_bytes` with no context), install the strict +session as the driver's *sender context*: + +.. code-block:: python + + from datafusion import SessionContext + from datafusion.ipc import set_sender_ctx + + set_sender_ctx(SessionContext().with_python_udf_inlining(enabled=False)) + # Every subsequent pickle.dumps(expr) on this thread encodes + # without inlining the Python callable. + +Pair with a matching strict worker context +(:py:func:`~datafusion.ipc.set_worker_ctx`) so the ``pickle.loads`` +side also refuses inline payloads. Explicit +:py:meth:`Expr.to_bytes(ctx) ` and +:py:meth:`Expr.from_bytes(blob, ctx=ctx) ` calls +honor the supplied ``ctx`` directly and ignore the sender / worker +contexts. + +The toggle only narrows the :py:meth:`Expr.from_bytes` surface; +:py:func:`pickle.loads` on untrusted bytes remains unsafe regardless +of this setting. See the `Security`_ section below for the full +threat model. + +Security +~~~~~~~~ + +.. warning:: + + Reconstructing an expression containing a Python UDF executes + arbitrary Python code on the receiver — pickle is doing the work + under the hood and pickle is unsafe on untrusted input (see the + `pickle module security warning + `_ + in the Python standard library docs). Only accept expressions + from trusted sources. For untrusted-source workflows, disable + Python UDF inlining (see above), restrict senders to built-in + functions and pre-registered Rust-side UDFs, and avoid + :py:func:`pickle.loads` on externally supplied bytes entirely. + +Reference: session context slots +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +There is only one type — :py:class:`SessionContext`. It can occupy +up to four *slots* in a running program: + +.. list-table:: + :header-rows: 1 + :widths: 12 18 40 30 + + * - Slot + - Lifetime + - Purpose + - Set how + * - User-held + - Local variable / attribute + - Build and run queries + - ``ctx = SessionContext(...)`` + * - Global + - Process singleton (lazy-init) + - Backs module-level + :py:func:`~datafusion.io.read_parquet`, + :py:func:`~datafusion.io.read_csv`, + :py:func:`~datafusion.io.read_json`, + :py:func:`~datafusion.io.read_avro`; final fallback for + :py:meth:`Expr.from_bytes` + - Implicit; access via + :py:meth:`SessionContext.global_ctx` + * - Sender + - Thread-local on the driver + - Codec settings for outbound :py:func:`pickle.dumps` / + :py:meth:`Expr.to_bytes` without ``ctx`` + - :py:func:`~datafusion.ipc.set_sender_ctx` + * - Worker + - Thread-local on the worker + - Function registry for inbound :py:func:`pickle.loads` / + :py:meth:`Expr.from_bytes` without ``ctx`` + - :py:func:`~datafusion.ipc.set_worker_ctx` + +The same :py:class:`SessionContext` object may occupy more than one +slot simultaneously — installing it into a slot is a reference, not +a copy. A non-distributed program only ever uses the user-held slot; +the global slot is invisible unless you call top-level ``read_*`` +helpers. + +Resolution order on the worker side is *explicit argument → +worker context → global context.* Explicit ``ctx=`` on +:py:meth:`Expr.from_bytes` always wins; the sender slot is ignored +on decode and the worker slot is ignored on encode. + +Sharp edges: + +* Sender and worker slots are **thread-local**. Background threads + on either side see ``None`` until they install their own. +* Under the ``fork`` start method, the parent's ``threading.local()`` + values are copied into the child by copy-on-write — a forked + worker initially observes whatever sender / worker slot the parent + had set, until the worker writes its own value (or calls the + matching ``clear_*_ctx``). ``spawn`` and ``forkserver`` workers + start with empty thread-local slots. Treat the slot as + uninitialized on worker entry and install (or clear) it explicitly + in the worker initializer; do not rely on inherited state. +* The global slot persists across ``fork`` workers (copy-on-write + memory inherit) but not across ``spawn`` / ``forkserver`` workers + (fresh process — register or install a worker context on + start-up). +* The inlining toggle is per-context state, not a global switch. + Two contexts with different toggles can coexist in one process. + +Query-level distribution via datafusion-distributed +--------------------------------------------------- + +🚧 *Work in progress upstream — not yet usable from datafusion-python.* + +`datafusion-distributed `_ +splits a single physical plan into stages and runs each stage on a +different worker node. The driver writes a SQL or DataFrame query +once; the runtime handles partitioning, shuffles, and reassembly. + +A datafusion-python integration is in development. This section will +document the integration once it lands. In the meantime, the +expression-level approach above covers most use cases that do not +require automatic plan partitioning. + +Query-level distribution via Apache Ballista +-------------------------------------------- + +🚧 *Work in progress upstream — not yet usable from datafusion-python.* + +`Apache Ballista `_ +provides distributed query execution on top of DataFusion with a +scheduler / executor model better suited to long-lived cluster +deployments. A datafusion-python integration is on the roadmap; this +section will fill in once the integration is usable. + +See also +-------- + +* :py:mod:`datafusion.ipc` — worker context API. +* ``examples/multiprocessing_pickle_expr.py`` — runnable + ``multiprocessing.Pool`` example that ships a different parametric + expression to each worker and collects results back. +* ``examples/ray_pickle_expr.py`` — runnable Ray actor example. diff --git a/examples/README.md b/examples/README.md index 3024c782f..e0e3056d9 100644 --- a/examples/README.md +++ b/examples/README.md @@ -44,6 +44,11 @@ Here is a direct link to the file used in the examples: - [Register a Python UDF with DataFusion](./python-udf.py) - [Register a Python UDAF with DataFusion](./python-udaf.py) +### Distributing DataFusion expressions + +- [Fan out distinct expressions to a multiprocessing pool](./multiprocessing_pickle_expr.py) +- [Distribute expression evaluation across Ray actors](./ray_pickle_expr.py) + ### Substrait Support - [Serialize query plans using Substrait](./substrait.py) diff --git a/examples/datafusion-ffi-example/python/tests/_test_pickle_strict_ffi.py b/examples/datafusion-ffi-example/python/tests/_test_pickle_strict_ffi.py new file mode 100644 index 000000000..67c0b245a --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_pickle_strict_ffi.py @@ -0,0 +1,127 @@ +# 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. + +"""Strict-mode Expr round-trip with an FFI-capsule scalar UDF. + +Verifies the by-name path: an FFI-imported UDF (no +``PythonFunctionScalarUDF`` downcast on the codec) serializes by name +and resolves from the receiver's function registry on decode. Covers +both the explicit ``Expr.to_bytes(ctx)`` / ``Expr.from_bytes(ctx=...)`` +API and the ``pickle.dumps`` / ``pickle.loads`` route through the +sender / worker context slots. +""" + +from __future__ import annotations + +import pickle + +import pyarrow as pa +import pytest +from datafusion import Expr, SessionContext, col, udf +from datafusion.ipc import ( + clear_sender_ctx, + clear_worker_ctx, + set_sender_ctx, + set_worker_ctx, +) +from datafusion_ffi_example import IsNullUDF + + +@pytest.fixture(autouse=True) +def _reset_thread_locals(): + """Ensure no sender / worker context leaks across tests.""" + clear_worker_ctx() + clear_sender_ctx() + yield + clear_worker_ctx() + clear_sender_ctx() + + +def _strict_session_with_ffi_udf(): + """Build a strict-mode session with the FFI ``IsNullUDF`` registered.""" + ctx = SessionContext().with_python_udf_inlining(enabled=False) + my_udf = udf(IsNullUDF()) + ctx.register_udf(my_udf) + return ctx, my_udf + + +def test_strict_ffi_udf_expr_roundtrip_via_to_bytes(): + """Strict-mode encode emits a by-name payload; receiver resolves + ``my_custom_is_null`` from its registered functions and the decoded + expression evaluates to the same result as the original.""" + sender, my_udf = _strict_session_with_ffi_udf() + receiver, _ = _strict_session_with_ffi_udf() + + expr = my_udf(col("a")) + blob = expr.to_bytes(sender) + restored = Expr.from_bytes(blob, ctx=receiver) + + assert "my_custom_is_null" in restored.canonical_name() + + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, None, 4], type=pa.int64())], names=["a"] + ) + receiver.register_record_batches("t", [[batch]]) + out = receiver.table("t").select(restored.alias("r")).collect() + expected = pa.array([False, False, True, False], type=pa.bool_()) + assert out[0].column(0) == expected + + +def test_strict_ffi_udf_pickle_roundtrip_via_thread_locals(): + """Driver installs a strict sender context; worker installs a + matching strict receiver. ``pickle.dumps`` / ``pickle.loads`` route + through them and the FFI UDF resolves by name on decode.""" + sender, my_udf = _strict_session_with_ffi_udf() + receiver, _ = _strict_session_with_ffi_udf() + + expr = my_udf(col("a")) + + set_sender_ctx(sender) + try: + blob = pickle.dumps(expr) + finally: + clear_sender_ctx() + + set_worker_ctx(receiver) + try: + restored = pickle.loads(blob) # noqa: S301 + finally: + clear_worker_ctx() + + assert "my_custom_is_null" in restored.canonical_name() + + +def test_strict_ffi_udf_smaller_than_inline_python_udf(): + """Sanity-check the wire size claim: strict-mode FFI UDF bytes are + a small by-name payload, dramatically smaller than the inline form + of a Python UDF with the same arity. Confirms the encode path + actually took the by-name branch instead of falling through to an + inline path.""" + sender, my_udf = _strict_session_with_ffi_udf() + ffi_blob = my_udf(col("a")).to_bytes(sender) + + inline_ctx = SessionContext() + py_udf = udf( + lambda arr: pa.array([v.as_py() is None for v in arr]), + [pa.int64()], + pa.bool_(), + volatility="immutable", + name="py_is_null", + ) + py_blob = py_udf(col("a")).to_bytes(inline_ctx) + + assert len(ffi_blob) < len(py_blob) // 4 diff --git a/examples/multiprocessing_pickle_expr.py b/examples/multiprocessing_pickle_expr.py new file mode 100644 index 000000000..73a99c2db --- /dev/null +++ b/examples/multiprocessing_pickle_expr.py @@ -0,0 +1,172 @@ +# 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. + +"""Distribute different DataFusion expressions to worker processes. + +For background — the shipped-expression model, what travels inline vs +by name, portability requirements, and the security threat model — +see ``docs/source/user-guide/distributing-work.rst``. + +Builds a list of parametric expressions in the driver — each closing +over a different threshold value — ships one per worker via +``multiprocessing.Pool``, and collects the results back. The closure +state forces the cloudpickle path (a by-name registration would lose +the captured threshold), so this is a real test of the expression- +pickling story rather than a same-expression fan-out. + +Worker layout: + +* Each worker receives a different ``(label, expr)`` task. +* Each worker materializes the shared dataset locally and runs its + own expression against it. +* The result and the worker's PID travel back to the driver, so the + output makes it visible that the work was spread across processes. + +Run: + python examples/multiprocessing_pickle_expr.py +""" + +from __future__ import annotations + +import multiprocessing as mp +import os + +import pyarrow as pa +from datafusion import Expr, SessionContext, col, udaf, udf +from datafusion import functions as F +from datafusion.user_defined import Accumulator, AggregateUDF, ScalarUDF + +# A shared input dataset. In a production pipeline this would live on +# object storage; here we hand-roll a small batch so the example runs +# without any I/O setup. +DATASET = { + "value": [3, 17, 42, 5, 88, 21, 9, 56, 4, 73, 12, 31], +} + + +def make_above_threshold_udf(threshold: int) -> ScalarUDF: + """Build a scalar UDF that returns 1 where ``value > threshold`` else 0. + + The threshold is captured in the closure, so cloudpickle has to + walk into the function body to ship the value across processes — + a by-name registration on the worker would collapse every + threshold into the same callable and lose the per-task state. + """ + + def above(arr: pa.Array) -> pa.Array: + # `v.as_py() or 0` coerces nulls to 0 — the demo dataset has no + # nulls, but real-world code should decide explicitly how nulls + # compare against the threshold. + return pa.array([1 if (v.as_py() or 0) > threshold else 0 for v in arr]) + + return udf( + above, + [pa.int64()], + pa.int64(), + volatility="immutable", + name=f"above_{threshold}", + ) + + +class _SumAccumulator(Accumulator): + """Tiny aggregate UDF state used to demonstrate UDAFs travel too.""" + + def __init__(self) -> None: + self._total = 0 + + def state(self) -> list[pa.Scalar]: + return [pa.scalar(self._total, type=pa.int64())] + + def update(self, values: pa.Array) -> None: + for v in values: + self._total += v.as_py() or 0 + + def merge(self, states: list[pa.Array]) -> None: + for s in states: + self._total += s[0].as_py() + + def evaluate(self) -> pa.Scalar: + return pa.scalar(self._total, type=pa.int64()) + + +def _build_sum_udaf() -> AggregateUDF: + return udaf( + _SumAccumulator, + [pa.int64()], + pa.int64(), + [pa.int64()], + "immutable", + name="my_sum", + ) + + +def evaluate_in_worker(task: tuple[str, Expr]) -> tuple[str, int, int]: + """Run one expression against the shared dataset. + + ``task`` arrived here via the pool's automatic pickling. The Python + callable inside the expression (including its captured threshold) + was reconstructed by the codec — the worker did not have to + register anything before this call. + """ + label, expr = task + ctx = SessionContext() + df = ctx.from_pydict(DATASET) + # ``expr`` is an aggregate over the whole batch; ``aggregate`` keeps + # a single row of output, which we read as a Python int. + result_df = df.aggregate([], [expr.alias("result")]) + result = result_df.to_pydict()["result"][0] + return label, result, os.getpid() + + +def build_tasks() -> list[tuple[str, Expr]]: + """Return ``(label, expr)`` pairs — one task per worker invocation. + + Mixes scalar-UDF-in-aggregate and pure-aggregate work to show both + UDF kinds round-tripping through pickle. + """ + sum_udaf = _build_sum_udaf() + tasks: list[tuple[str, Expr]] = [] + + # Three "count values strictly above threshold T" tasks built from + # closure-capturing scalar UDFs. + for threshold in (10, 30, 60): + above_udf = make_above_threshold_udf(threshold) + tasks.append((f"count_above_{threshold}", F.sum(above_udf(col("value"))))) + + # One pure aggregate UDF task. + tasks.append(("custom_sum", sum_udaf(col("value")))) + + return tasks + + +def main() -> None: + tasks = build_tasks() + + # ``forkserver`` works on every POSIX platform and is the Python 3.14 + # default for POSIX. ``spawn`` would also work; ``fork`` is unsafe + # with pyarrow/tokio on macOS. + mp_ctx = mp.get_context("forkserver") + with mp_ctx.Pool(processes=min(4, len(tasks))) as pool: + results = pool.map(evaluate_in_worker, tasks) + + print(f"driver pid: {os.getpid()}") + for label, value, worker_pid in results: + print(f" [{label:>16}] = {value:>6} (worker pid: {worker_pid})") + + +if __name__ == "__main__": + main() diff --git a/examples/ray_pickle_expr.py b/examples/ray_pickle_expr.py new file mode 100644 index 000000000..04cea463d --- /dev/null +++ b/examples/ray_pickle_expr.py @@ -0,0 +1,86 @@ +# 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. + +"""Distribute DataFusion expressions to Ray actors. + +For background — the shipped-expression model, what travels inline vs +by name, portability requirements, and the security threat model — +see ``docs/source/user-guide/distributing-work.rst``. + +Build an expression in the driver, ship it to a pool of Ray actors, and +have each actor evaluate it against its own slice of data. Python UDFs +travel with the shipped expression — no actor-side registration needed. + +Prerequisites: + pip install ray + +Run: + python examples/ray_pickle_expr.py +""" + +import pyarrow as pa +import ray +from datafusion import Expr, SessionContext, col, lit, udf + + +def _build_double_udf(): + """Return the demo UDF used by the driver.""" + return udf( + lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]), + [pa.int64()], + pa.int64(), + volatility="immutable", + name="double", + ) + + +@ray.remote +class DataFusionWorker: + """A Ray actor with a private :class:`SessionContext`.""" + + def __init__(self) -> None: + self._ctx = SessionContext() + + def evaluate(self, expr: Expr, batch_pylist: list[int]) -> list[int]: + """Run the expression against an in-memory batch.""" + # `expr` arrived here via Ray's automatic argument serialization; + # the Python UDF inside it was reconstructed from the bytes — no + # pre-registration on this actor required. + df = self._ctx.from_pydict({"a": batch_pylist}) + out = df.with_column("result", expr).select("result") + return out.to_pydict()["result"] + + +def main() -> None: + ray.init(ignore_reinit_error=True) + + expr = _build_double_udf()(col("a")) + lit(1) + + workers = [DataFusionWorker.remote() for _ in range(2)] + batches = [[1, 2, 3], [10, 20, 30], [100, 200, 300]] + futures = [ + workers[i % len(workers)].evaluate.remote(expr, batch) + for i, batch in enumerate(batches) + ] + for batch, result in zip(batches, ray.get(futures), strict=True): + print(f"input {batch} -> {result}") + + ray.shutdown() + + +if __name__ == "__main__": + main() From 987228300b1c52215b4bb10a1cd5781c40648fbe Mon Sep 17 00:00:00 2001 From: kosiew Date: Fri, 29 May 2026 17:42:10 +0800 Subject: [PATCH 45/83] Export `to_datafusion_err` from the util crate root (#1487) * Re-export error helpers from crate root Publicly re-export curated error helpers, including to_datafusion_err, from the crate root. Add a regression test in crates/util/tests/root_exports.rs to ensure correct functionality in an integration-test context. * Restrict re-exported items in lib.rs Make only to_datafusion_err publicly re-exported from the crate root. Keep PyDataFusionError and PyDataFusionResult as private imports for internal use, enhancing encapsulation and reducing exposure of non-essential components. * feat: remove unnecessary blank line in lib.rs to improve code formatting * rm root_exports.rs --- crates/util/src/lib.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 72dc9aafc..75b8eeec4 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -36,9 +36,9 @@ use tokio::runtime::Runtime; use tokio::task::JoinHandle; use tokio::time::sleep; -use crate::errors::{PyDataFusionError, PyDataFusionResult, to_datafusion_err}; - pub mod errors; +pub use crate::errors::to_datafusion_err; +use crate::errors::{PyDataFusionError, PyDataFusionResult}; /// Utility to get the Tokio Runtime from Python #[inline] From 7df58e531bb949785cae3cca488a1ae55cb6d478 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 29 May 2026 08:09:42 -0400 Subject: [PATCH 46/83] feat: pass calling SessionContext to Python UDTF callbacks (#1555) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: pass calling SessionContext to Python UDTF callbacks DataFusion 53 added `TableFunctionImpl::call_with_args(TableFunctionArgs)` where `TableFunctionArgs` carries both the positional expression arguments and the calling `&dyn Session`. The pure-Python UDTF path previously discarded everything but the exprs. Thread the session through when the user callback's signature opts in by declaring a `session` keyword parameter (or `**kwargs`). At call time we downcast the `&dyn Session` to its canonical `SessionState` impl and build a fresh `SessionContext` over the same Arc-shared state, exposed to Python as a `datafusion.SessionContext` wrapper. Existing callbacks whose signatures do not declare `session` continue to be called with the positional expression arguments only — no behavior change for current users. Note: a UDTF body cannot drive a fresh `ctx.sql(...).collect()` on the passed-in session because the outer SQL execution already holds the tokio runtime. Use the session for metadata access (catalogs, UDF lookups, config) rather than nested DataFrame collection. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: clarify py_session_from_session downcast is defensive The doc comment implied a foreign FFI session was a real input. No current path reaches a pure-Python UDTF with a non-SessionState session: the SQL planner and __call__ both hand a SessionState, and a ForeignSession would only arrive via FFI-export of the UDTF, which datafusion-python does not do. Reword to state the guard is defensive and rewrap the error string. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: opt-in UDTF session injection via with_session flag Replaces signature sniffing with an explicit ``with_session=True`` kwarg on ``TableFunction`` / ``udtf``. Avoids name-based detection footguns (positional-only ``session`` params, accidental ``**kwargs`` opt-in, shadowing by unrelated params) and makes author intent visible at registration. Also documents the feature in the UDTF user guide. Rust field renamed ``accepts_session`` -> ``inject_session_on_call`` to match the Python-side opt-in semantics. Co-Authored-By: Claude Opus 4.7 (1M context) * fix: reject with_session=True for FFI UDTFs and qualify mutation docs Raise TypeError when with_session=True is combined with an FFI-exported table function (one exposing __datafusion_table_function__). The Rust FFI branch does not consult the flag, so it would silently be dropped; guard both TableFunction.__init__ and the udtf() convenience entry. Qualify the doc claim that mutations through the injected session propagate to the caller: registry mutations do (shared Arc registries), but config changes do not (SessionConfig is cloned). Mirror the caveat in TableFunction.__init__ per the user-guide caveats convention. Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/udtf.rs | 94 +++++++++++++--- .../common-operations/udf-and-udfa.rst | 39 +++++++ python/datafusion/user_defined.py | 100 +++++++++++++++--- python/tests/test_udtf.py | 99 +++++++++++++++++ 4 files changed, 303 insertions(+), 29 deletions(-) diff --git a/crates/core/src/udtf.rs b/crates/core/src/udtf.rs index b3de25e52..cffa0c12a 100644 --- a/crates/core/src/udtf.rs +++ b/crates/core/src/udtf.rs @@ -18,20 +18,34 @@ use std::ptr::NonNull; use std::sync::Arc; -use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider}; -use datafusion::error::Result as DataFusionResult; +use datafusion::catalog::{Session, TableFunctionArgs, TableFunctionImpl, TableProvider}; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::SessionContext; +use datafusion::execution::session_state::SessionState; use datafusion::logical_expr::Expr; use datafusion_ffi::udtf::FFI_TableFunction; use pyo3::IntoPyObjectExt; use pyo3::exceptions::{PyImportError, PyTypeError}; use pyo3::prelude::*; -use pyo3::types::{PyCapsule, PyTuple, PyType}; +use pyo3::types::{PyCapsule, PyDict, PyTuple, PyType}; use crate::context::PySessionContext; use crate::errors::{py_datafusion_err, to_datafusion_err}; use crate::expr::PyExpr; use crate::table::PyTable; +/// A pure-Python UDTF callable plus the metadata we discovered about it +/// at registration time. +#[derive(Debug, Clone)] +pub(crate) struct PythonTableFunctionCallable { + pub(crate) callable: Arc>, + /// When true, the calling :class:`SessionContext` is passed to the + /// callable as a ``session`` keyword argument on every invocation. + /// Opt-in at registration time via ``with_session=True`` on the + /// Python wrapper. + pub(crate) inject_session_on_call: bool, +} + /// Represents a user defined table function #[pyclass(from_py_object, frozen, name = "TableFunction", module = "datafusion")] #[derive(Debug, Clone)] @@ -40,21 +54,21 @@ pub struct PyTableFunction { pub(crate) inner: PyTableFunctionInner, } -// TODO: Implement pure python based user defined table functions #[derive(Debug, Clone)] pub(crate) enum PyTableFunctionInner { - PythonFunction(Arc>), + PythonFunction(PythonTableFunctionCallable), FFIFunction(Arc), } #[pymethods] impl PyTableFunction { #[new] - #[pyo3(signature=(name, func, session))] + #[pyo3(signature=(name, func, session, inject_session_on_call=false))] pub fn new( name: &str, func: Bound<'_, PyAny>, session: Option>, + inject_session_on_call: bool, ) -> PyResult { let inner = if func.hasattr("__datafusion_table_function__")? { let py = func.py(); @@ -80,8 +94,10 @@ impl PyTableFunction { PyTableFunctionInner::FFIFunction(foreign_func) } else { - let py_obj = Arc::new(func.unbind()); - PyTableFunctionInner::PythonFunction(py_obj) + PyTableFunctionInner::PythonFunction(PythonTableFunctionCallable { + callable: Arc::new(func.unbind()), + inject_session_on_call, + }) }; Ok(Self { @@ -107,20 +123,66 @@ impl PyTableFunction { } } +/// Materialize a fresh :class:`PySessionContext` from the borrowed +/// ``&dyn Session`` handed in at call time. +/// +/// Upstream invokes ``call_with_args`` with a trait-object reference +/// rather than an owned context; we downcast it to the canonical +/// :class:`SessionState` impl and rebuild a :class:`SessionContext` +/// (sharing the same registries via the Arc-heavy interior of +/// :class:`SessionState`). +/// +/// The downcast is defensive. Every path that reaches a pure-Python +/// UDTF today hands us a `SessionState`: the SQL planner builds the +/// args from its own `SessionState`, and `PyTableFunction::__call__` +/// uses the global context's state. A non-`SessionState` session +/// (e.g. a `ForeignSession`) would only arrive if this UDTF were +/// exported across the FFI boundary to a foreign-library consumer, +/// which datafusion-python does not do. Should that change, this +/// returns an error rather than silently misbehaving. +fn py_session_from_session(session: &dyn Session) -> DataFusionResult { + let state = session + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution( + "Cannot expose this UDTF's calling session to Python: the \ + session is not a SessionState. Drop the `session` keyword \ + from the callback signature to fall back to the \ + expression-only call form." + .to_string(), + ) + })?; + Ok(PySessionContext::from(SessionContext::new_with_state( + state.clone(), + ))) +} + #[allow(clippy::result_large_err)] fn call_python_table_function( - func: &Arc>, - args: &[Expr], + func: &PythonTableFunctionCallable, + args: TableFunctionArgs, ) -> DataFusionResult> { - let args = args + let py_session = if func.inject_session_on_call { + Some(py_session_from_session(args.session())?) + } else { + None + }; + let py_exprs = args + .exprs() .iter() .map(|arg| PyExpr::from(arg.clone())) .collect::>(); - // move |args: &[ArrayRef]| -> Result { Python::attach(|py| { - let py_args = PyTuple::new(py, args)?; - let provider_obj = func.call1(py, py_args)?; + let py_args = PyTuple::new(py, py_exprs)?; + let provider_obj = if let Some(session) = py_session { + let kwargs = PyDict::new(py); + kwargs.set_item("session", session.into_pyobject(py)?)?; + func.callable.call(py, py_args, Some(&kwargs))? + } else { + func.callable.call1(py, py_args)? + }; let provider = provider_obj.bind(py).clone(); Ok::, PyErr>(PyTable::new(provider, None)?.table) @@ -132,8 +194,8 @@ impl TableFunctionImpl for PyTableFunction { fn call_with_args(&self, args: TableFunctionArgs) -> DataFusionResult> { match &self.inner { PyTableFunctionInner::FFIFunction(func) => func.call_with_args(args), - PyTableFunctionInner::PythonFunction(obj) => { - call_python_table_function(obj, args.exprs()) + PyTableFunctionInner::PythonFunction(callable) => { + call_python_table_function(callable, args) } } } 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 59c47b595..918c2e29e 100644 --- a/docs/source/user-guide/common-operations/udf-and-udfa.rst +++ b/docs/source/user-guide/common-operations/udf-and-udfa.rst @@ -431,3 +431,42 @@ that you wish to expose via PyO3, you need to expose it as a ``PyCapsule``. PyCapsule::new(py, provider, Some(name)) } } + +Accessing the Calling Session +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Pure-Python UDTFs can opt into receiving the calling +:py:class:`~datafusion.SessionContext` by registering with +``with_session=True``. The context is passed as a ``session`` keyword +argument on every invocation. Use it to look up registered tables, +UDFs, or session configuration from inside the callback. + +.. code-block:: python + + from datafusion import SessionContext, Table, udtf + from datafusion.context import TableProviderExportable + import pyarrow as pa + import pyarrow.dataset as ds + + @udtf("list_tables", with_session=True) + def list_tables(*, session: SessionContext) -> TableProviderExportable: + names = sorted(session.catalog().schema().names()) + batch = pa.RecordBatch.from_pydict({"name": names}) + return Table(ds.dataset([batch])) + + ctx = SessionContext() + ctx.register_batch("t1", pa.RecordBatch.from_pydict({"x": [1]})) + ctx.register_udtf(list_tables) + ctx.sql("SELECT * FROM list_tables()").show() + +Without ``with_session=True``, the callback receives only the positional +expression arguments. The flag is opt-in so existing UDTFs keep working +unchanged. + +The injected ``session`` is a fresh :py:class:`~datafusion.SessionContext` +wrapper backed by the same underlying state as the caller, so registries +(tables, UDFs, catalogs) are visible. Registry mutations (e.g. registering +a new table or UDF) propagate to the live session because the registries +are reference-counted and shared. Configuration changes made through the +wrapper (e.g. setting session options) do **not** propagate — the wrapper +holds its own clone of the session config. diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index ced707a96..81a516af8 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -1102,6 +1102,24 @@ def from_pycapsule(func: WindowUDFExportable) -> WindowUDF: ) +def _wrap_session_kwarg_for_udtf(func: Callable[..., Any]) -> Callable[..., Any]: + """Adapt the raw internal session pyo3 object back to a Python wrapper. + + The Rust call site forwards a ``datafusion._internal.SessionContext``, + but UDTF authors expect to interact with the public + :class:`datafusion.SessionContext` wrapper. This closure wraps the + internal object once per call before delegating to ``func``. + """ + + @functools.wraps(func, updated=()) + def adapter(*args: Any, session: Any, **kwargs: Any) -> Any: + wrapped = SessionContext.__new__(SessionContext) + wrapped.ctx = session + return func(*args, session=wrapped, **kwargs) + + return adapter + + class TableFunction: """Class for performing user-defined table functions (UDTF). @@ -1110,14 +1128,44 @@ class TableFunction: """ def __init__( - self, name: str, func: Callable[[], any], ctx: SessionContext | None = None + self, + name: str, + func: Callable[..., Any], + ctx: SessionContext | None = None, + *, + with_session: bool = False, ) -> None: """Instantiate a user-defined table function (UDTF). + Set ``with_session=True`` to have the calling + :class:`SessionContext` passed as a ``session`` keyword argument + on each invocation. Use it inside the callback to look up + registered tables, UDFs, or session configuration. When + ``with_session`` is ``False`` (the default), ``func`` is invoked + with the positional expression arguments only. + + ``with_session=True`` is only supported for pure-Python callables. + Passing it together with an FFI-exported table function (one + exposing ``__datafusion_table_function__``) raises + :class:`TypeError`. + + Registry mutations performed through the injected session (such + as registering tables or UDFs) propagate to the caller's + :class:`SessionContext` because the registries are shared. + Configuration changes do **not** propagate; the wrapper holds + its own clone of the session config. + See :py:func:`udtf` for a convenience function and argument descriptions. """ - self._udtf = df_internal.TableFunction(name, func, ctx) + if with_session and hasattr(func, "__datafusion_table_function__"): + msg = ( + "`with_session=True` is not supported for FFI-exported table " + "functions; session injection requires a pure-Python callable." + ) + raise TypeError(msg) + registered = _wrap_session_kwarg_for_udtf(func) if with_session else func + self._udtf = df_internal.TableFunction(name, registered, ctx, with_session) def __call__(self, *args: Expr) -> Any: """Execute the UDTF and return a table provider.""" @@ -1128,47 +1176,73 @@ def __call__(self, *args: Expr) -> Any: @staticmethod def udtf( name: str, + *, + with_session: bool = False, ) -> Callable[..., Any]: ... @overload @staticmethod def udtf( - func: Callable[[], Any], + func: Callable[..., Any], name: str, + *, + with_session: bool = False, ) -> TableFunction: ... @staticmethod - def udtf(*args: Any, **kwargs: Any): - """Create a new User-Defined Table Function (UDTF).""" + def udtf(*args: Any, with_session: bool = False, **kwargs: Any): + """Create a new User-Defined Table Function (UDTF). + + Pass ``with_session=True`` to have the calling + :class:`SessionContext` injected as a ``session`` keyword + argument on each invocation. + """ if args and callable(args[0]): # Case 1: Used as a function, require the first parameter to be callable - return TableFunction._create_table_udf(*args, **kwargs) + return TableFunction._create_table_udf( + *args, with_session=with_session, **kwargs + ) if args and hasattr(args[0], "__datafusion_table_function__"): # Case 2: We have a datafusion FFI provided function + if with_session: + msg = ( + "`with_session=True` is not supported for FFI-exported " + "table functions; session injection requires a " + "pure-Python callable." + ) + raise TypeError(msg) return TableFunction(args[1], args[0]) # Case 3: Used as a decorator with parameters - return TableFunction._create_table_udf_decorator(*args, **kwargs) + return TableFunction._create_table_udf_decorator( + *args, with_session=with_session, **kwargs + ) @staticmethod def _create_table_udf( func: Callable[..., Any], name: str, + *, + with_session: bool = False, ) -> TableFunction: """Create a TableFunction instance from function arguments.""" if not callable(func): msg = "`func` must be callable." raise TypeError(msg) - return TableFunction(name, func) + return TableFunction(name, func, with_session=with_session) @staticmethod def _create_table_udf_decorator( name: str | None = None, - ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]: - """Create a decorator for a WindowUDF.""" - - def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]: - return TableFunction._create_table_udf(func, name) + *, + with_session: bool = False, + ) -> Callable[[Callable[..., Any]], TableFunction]: + """Create a decorator for a TableFunction.""" + + def decorator(func: Callable[..., Any]) -> TableFunction: + return TableFunction._create_table_udf( + func, name, with_session=with_session + ) return decorator diff --git a/python/tests/test_udtf.py b/python/tests/test_udtf.py index 925a8ba01..dcb2bacc3 100644 --- a/python/tests/test_udtf.py +++ b/python/tests/test_udtf.py @@ -17,8 +17,10 @@ import pyarrow as pa import pyarrow.dataset as ds +import pytest from datafusion import Expr, SessionContext, Table, udtf from datafusion.context import TableProviderExportable +from datafusion.user_defined import TableFunction def python_table_function_inner( @@ -134,3 +136,100 @@ def string_arg_func(prefix: Expr) -> TableProviderExportable: result = ctx.sql("SELECT * FROM string_arg_func('test')").collect() assert len(result) == 1 assert result[0].schema.names == ["test_a", "test_b"] + + +def test_python_table_function_receives_session() -> None: + """A UDTF registered ``with_session=True`` gets the calling ctx.""" + ctx = SessionContext() + captured: list[SessionContext] = [] + + @udtf("session_aware_func", with_session=True) + def session_aware_func(*, session: SessionContext) -> TableProviderExportable: + captured.append(session) + batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3]}) + return Table(ds.dataset([batch])) + + ctx.register_udtf(session_aware_func) + result = ctx.sql("SELECT * FROM session_aware_func()").collect() + + assert len(captured) == 1 + assert isinstance(captured[0], SessionContext) + # Sharing the same catalog confirms the wrapper points at the caller's state. + assert captured[0].catalog().schema().names() == ctx.catalog().schema().names() + assert result[0].column(0).to_pylist() == [1, 2, 3] + + +def test_python_table_function_session_used_for_metadata() -> None: + """The UDTF can inspect session state through the passed-in context.""" + ctx = SessionContext() + base_batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) + ctx.register_batch("base_tbl", base_batch) + + seen_tables: list[set[str]] = [] + + @udtf("table_inventory", with_session=True) + def table_inventory(*, session: SessionContext) -> TableProviderExportable: + # Stash the visible tables to verify the session wired through. + seen_tables.append(session.catalog().schema().names()) + batch = pa.RecordBatch.from_pydict({"name": ["base_tbl"]}) + return Table(ds.dataset([batch])) + + ctx.register_udtf(table_inventory) + result = ctx.sql("SELECT * FROM table_inventory()").collect() + + assert seen_tables == [{"base_tbl"}] + assert result[0].column(0).to_pylist() == ["base_tbl"] + + +def test_python_table_function_class_callable_with_session() -> None: + """Class-based UDTFs opt in via ``with_session=True``.""" + ctx = SessionContext() + captured: list[SessionContext] = [] + + class SessionAware: + def __call__( + self, n: Expr, *, session: SessionContext + ) -> TableProviderExportable: + captured.append(session) + count = n.to_variant().value_i64() + batch = pa.RecordBatch.from_pydict({"a": list(range(count))}) + return Table(ds.dataset([batch])) + + ctx.register_udtf(udtf(SessionAware(), "session_class_func", with_session=True)) + result = ctx.sql("SELECT * FROM session_class_func(3)").collect() + + assert len(captured) == 1 + assert isinstance(captured[0], SessionContext) + assert result[0].column(0).to_pylist() == [0, 1, 2] + + +def test_python_table_function_without_session_flag_no_injection() -> None: + """Default registration (no ``with_session``) calls func positionally.""" + ctx = SessionContext() + + @udtf("plain_func") + def plain_func(n: Expr) -> TableProviderExportable: + count = n.to_variant().value_i64() + batch = pa.RecordBatch.from_pydict({"a": list(range(count))}) + return Table(ds.dataset([batch])) + + ctx.register_udtf(plain_func) + result = ctx.sql("SELECT * FROM plain_func(4)").collect() + + assert result[0].column(0).to_pylist() == [0, 1, 2, 3] + + +def test_with_session_rejected_for_ffi_table_function() -> None: + """`with_session=True` is incompatible with FFI-exported table functions.""" + + class FakeFFITableFunction: + # Presence of this attribute is what marks a function as FFI-exported. + __datafusion_table_function__ = "stub" + + fake = FakeFFITableFunction() + + with pytest.raises(TypeError, match="FFI-exported table functions"): + udtf(fake, "fake_ffi", with_session=True) + + with pytest.raises(TypeError, match="FFI-exported table functions"): + TableFunction("fake_ffi", fake, with_session=True) From 744dd23ff3bfeeafbf532d71c909e0aedb7b2194 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 29 May 2026 09:25:16 -0400 Subject: [PATCH 47/83] chore: remove unused PyConfig (#1485) * Remove unused Config class that could not be connected to a SessionContext Closes #322. Co-Authored-By: Claude Opus 4.6 (1M context) * Minor correction after merge main --------- Co-authored-by: Claude Opus 4.6 (1M context) --- crates/core/src/config.rs | 104 ---------------------- crates/core/src/lib.rs | 3 - docs/source/contributor-guide/ffi.rst | 15 +--- docs/source/user-guide/upgrade-guides.rst | 27 ++++++ python/datafusion/__init__.py | 2 - python/tests/test_concurrency.py | 35 +------- python/tests/test_config.py | 42 --------- 7 files changed, 30 insertions(+), 198 deletions(-) delete mode 100644 crates/core/src/config.rs delete mode 100644 python/tests/test_config.py diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs deleted file mode 100644 index fdb693a12..000000000 --- a/crates/core/src/config.rs +++ /dev/null @@ -1,104 +0,0 @@ -// 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. - -use std::sync::Arc; - -use datafusion::config::ConfigOptions; -use parking_lot::RwLock; -use pyo3::prelude::*; -use pyo3::types::*; - -use crate::common::data_type::PyScalarValue; -use crate::errors::PyDataFusionResult; -#[pyclass( - from_py_object, - name = "Config", - module = "datafusion", - subclass, - frozen -)] -#[derive(Clone)] -pub(crate) struct PyConfig { - config: Arc>, -} - -#[pymethods] -impl PyConfig { - #[new] - fn py_new() -> Self { - Self { - config: Arc::new(RwLock::new(ConfigOptions::new())), - } - } - - /// Get configurations from environment variables - #[staticmethod] - pub fn from_env() -> PyDataFusionResult { - Ok(Self { - config: Arc::new(RwLock::new(ConfigOptions::from_env()?)), - }) - } - - /// Get a configuration option - pub fn get<'py>(&self, key: &str, py: Python<'py>) -> PyResult> { - let value: Option> = { - let options = self.config.read(); - options - .entries() - .into_iter() - .find_map(|entry| (entry.key == key).then_some(entry.value.clone())) - }; - - match value { - Some(value) => Ok(value.into_pyobject(py)?), - None => Ok(None::.into_pyobject(py)?), - } - } - - /// Set a configuration option - pub fn set(&self, key: &str, value: Py, py: Python) -> PyDataFusionResult<()> { - let scalar_value: PyScalarValue = value.extract(py)?; - let mut options = self.config.write(); - options.set(key, scalar_value.0.to_string().as_str())?; - Ok(()) - } - - /// Get all configuration options - pub fn get_all(&self, py: Python) -> PyResult> { - let entries: Vec<(String, Option)> = { - let options = self.config.read(); - options - .entries() - .into_iter() - .map(|entry| (entry.key.clone(), entry.value.clone())) - .collect() - }; - - let dict = PyDict::new(py); - for (key, value) in entries { - dict.set_item(key, value.into_pyobject(py)?)?; - } - Ok(dict.into()) - } - - fn __repr__(&self, py: Python) -> PyResult { - match self.get_all(py) { - Ok(result) => Ok(format!("Config({result})")), - Err(err) => Ok(format!("Error: {:?}", err.to_string())), - } - } -} diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 8b622d344..79bf77717 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -31,8 +31,6 @@ pub mod catalog; pub mod codec; pub mod common; -#[allow(clippy::borrow_deref_ref)] -mod config; #[allow(clippy::borrow_deref_ref)] pub mod context; #[allow(clippy::borrow_deref_ref)] @@ -97,7 +95,6 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/docs/source/contributor-guide/ffi.rst b/docs/source/contributor-guide/ffi.rst index e0158e0a2..c89b99849 100644 --- a/docs/source/contributor-guide/ffi.rst +++ b/docs/source/contributor-guide/ffi.rst @@ -149,20 +149,7 @@ interior-mutable state. In practice this means that any ``#[pyclass]`` containin ``Arc>`` or similar synchronized primitive must opt into ``#[pyclass(frozen)]`` unless there is a compelling reason not to. -The :mod:`datafusion` configuration helpers illustrate the preferred pattern. The -``PyConfig`` class in :file:`src/config.rs` stores an ``Arc>`` and is -explicitly frozen so callers interact with configuration state through provided methods -instead of mutating the container directly: - -.. code-block:: rust - - #[pyclass(from_py_object, name = "Config", module = "datafusion", subclass, frozen)] - #[derive(Clone)] - pub(crate) struct PyConfig { - config: Arc>, - } - -The same approach applies to execution contexts. ``PySessionContext`` in +The execution context illustrates the preferred pattern. ``PySessionContext`` in :file:`src/context.rs` stays frozen even though it shares mutable state internally via ``SessionContext``. This ensures PyO3 tracks borrows correctly while Python-facing APIs clone the inner ``SessionContext`` or return new wrappers instead of mutating the diff --git a/docs/source/user-guide/upgrade-guides.rst b/docs/source/user-guide/upgrade-guides.rst index 2ac7f7703..12d0f1bb3 100644 --- a/docs/source/user-guide/upgrade-guides.rst +++ b/docs/source/user-guide/upgrade-guides.rst @@ -18,6 +18,33 @@ Upgrade Guides ============== +DataFusion 54.0.0 +----------------- + +The ``Config`` class has been removed. It was a standalone wrapper around +``ConfigOptions`` that could not be connected to a ``SessionContext``, making it +effectively unusable. Use :py:class:`~datafusion.context.SessionConfig` instead, +which is passed directly to ``SessionContext``. + +Before: + +.. code-block:: python + + from datafusion import Config + + config = Config() + config.set("datafusion.execution.batch_size", "4096") + # config could not be passed to SessionContext + +After: + +.. code-block:: python + + from datafusion import SessionConfig, SessionContext + + config = SessionConfig().set("datafusion.execution.batch_size", "4096") + ctx = SessionContext(config) + DataFusion 53.0.0 ----------------- diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index f8d523e0d..9c55f446c 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -68,7 +68,6 @@ from . import functions, ipc, object_store, substrait, unparser # The following imports are okay to remain as opaque to the user. -from ._internal import Config from .catalog import ( Catalog, Table, @@ -115,7 +114,6 @@ "Accumulator", "AggregateUDF", "Catalog", - "Config", "CsvReadOptions", "DFSchema", "DataFrame", diff --git a/python/tests/test_concurrency.py b/python/tests/test_concurrency.py index f790f9473..e7fe44f01 100644 --- a/python/tests/test_concurrency.py +++ b/python/tests/test_concurrency.py @@ -20,7 +20,7 @@ from concurrent.futures import ThreadPoolExecutor import pyarrow as pa -from datafusion import Config, SessionContext, col, lit +from datafusion import SessionContext, col, lit from datafusion import functions as f from datafusion.common import SqlSchema @@ -34,16 +34,14 @@ def _run_in_threads(fn, count: int = 8) -> None: def test_concurrent_access_to_shared_structures() -> None: - """Exercise SqlSchema, Config, and DataFrame concurrently.""" + """Exercise SqlSchema and DataFrame concurrently.""" schema = SqlSchema("concurrency") - config = Config() ctx = SessionContext() batch = pa.record_batch([pa.array([1, 2, 3], type=pa.int32())], names=["value"]) df = ctx.create_dataframe([[batch]]) - config_key = "datafusion.execution.batch_size" expected_rows = batch.num_rows def worker(index: int) -> None: @@ -54,41 +52,12 @@ def worker(index: int) -> None: assert isinstance(schema.views, list) assert isinstance(schema.functions, list) - config.set(config_key, str(1024 + index)) - assert config.get(config_key) is not None - # Access the full config map to stress lock usage. - assert config_key in config.get_all() - batches = df.collect() assert sum(batch.num_rows for batch in batches) == expected_rows _run_in_threads(worker, count=12) -def test_config_set_during_get_all() -> None: - """Ensure config writes proceed while another thread reads all entries.""" - - config = Config() - key = "datafusion.execution.batch_size" - - def reader() -> None: - for _ in range(200): - # get_all should not hold the lock while converting to Python objects - config.get_all() - - def writer() -> None: - for index in range(200): - config.set(key, str(1024 + index)) - - with ThreadPoolExecutor(max_workers=2) as executor: - reader_future = executor.submit(reader) - writer_future = executor.submit(writer) - reader_future.result(timeout=10) - writer_future.result(timeout=10) - - assert config.get(key) is not None - - def test_case_builder_reuse_from_multiple_threads() -> None: """Ensure the case builder can be safely reused across threads.""" diff --git a/python/tests/test_config.py b/python/tests/test_config.py deleted file mode 100644 index c1d7f97e1..000000000 --- a/python/tests/test_config.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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. - -import pytest -from datafusion import Config - - -@pytest.fixture -def config(): - return Config() - - -def test_get_then_set(config): - config_key = "datafusion.optimizer.filter_null_join_keys" - - assert config.get(config_key) == "false" - - config.set(config_key, "true") - assert config.get(config_key) == "true" - - -def test_get_all(config): - config_dict = config.get_all() - assert config_dict["datafusion.catalog.create_default_catalog_and_schema"] == "true" - - -def test_get_invalid_config(config): - assert config.get("not.valid.key") is None From 0840763851e47a40676dc47b29f5f3a62b0127d7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 29 May 2026 13:30:39 -0400 Subject: [PATCH 48/83] feat: accept distinct kwarg on sum and avg (#1556) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: accept distinct kwarg on sum and avg Upstream exposes `sum_distinct` / `avg_distinct` / `count_distinct` as sibling functions that call the same underlying UDAF with `distinct: bool = true`. The Rust binding side already routes `distinct=Some(true)` through the aggregate builder for `sum`, `avg`, and `count` — but only `count` exposed the kwarg on the Python wrapper. Add `distinct: bool = False` to `sum()` and `avg()` mirroring the existing `count()` signature, and update SKILL.md so the check-upstream audit does not re-flag the three upstream `*_distinct` shortcuts as gaps. The plan emitted by `sum(col, distinct=True)` matches what upstream's `sum_distinct(col)` builds. Co-Authored-By: Claude Opus 4.7 (1M context) * test: fold sum/avg distinct tests into parameterized aggregation test Move the standalone test_sum_distinct_kwarg and test_avg_distinct_kwarg from test_functions.py into the existing test_aggregation::test_aggregation parameterization, matching how distinct is already covered for median, array_agg, count, and bit_xor. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: clarify distinct kwarg on sum and avg Drop the unhelpful "upstream avg_distinct/sum_distinct shortcut" reference in favor of describing the actual behavior. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: note sum/avg distinct argument-order breaking change distinct is inserted before filter on sum and avg for consistency with the other aggregate functions, breaking positional filter callers. Add a DataFusion 54.0.0 upgrade-guide entry covering the migration. Co-Authored-By: Claude Opus 4.7 (1M context) * Update docs/source/user-guide/upgrade-guides.rst Co-authored-by: Nick <24689722+ntjohnson1@users.noreply.github.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Nick <24689722+ntjohnson1@users.noreply.github.com> --- .ai/skills/check-upstream/SKILL.md | 9 +++++++- docs/source/user-guide/upgrade-guides.rst | 20 ++++++++++++++++ python/datafusion/functions.py | 28 +++++++++++++++++++---- python/tests/test_aggregation.py | 2 ++ 4 files changed, 54 insertions(+), 5 deletions(-) diff --git a/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md index 23873feab..24b4e1bb1 100644 --- a/.ai/skills/check-upstream/SKILL.md +++ b/.ai/skills/check-upstream/SKILL.md @@ -88,11 +88,18 @@ The user may specify an area via `$ARGUMENTS`. If no area is specified or "all" - Python API: `python/datafusion/functions.py` (aggregate functions are mixed in with scalar functions) - Rust bindings: `crates/core/src/functions.rs` +**Evaluated and not requiring separate Python exposure:** +- `count_distinct` — covered by `count(expr, distinct=True)`. Both forms call + `count_udaf` with `distinct: bool = true` and produce the same logical plan. +- `sum_distinct` — covered by `sum(expr, distinct=True)`. +- `avg_distinct` — covered by `avg(expr, distinct=True)`. + **How to check:** 1. Fetch the upstream aggregate function documentation page 2. Compare against aggregate functions in `python/datafusion/functions.py` (check `__all__` list and function definitions) 3. A function is covered if it exists in the Python API, even if it aliases another function's Rust binding -4. Report only functions missing from the Python API +4. Check against the "evaluated and not requiring exposure" list before flagging as a gap +5. Report only functions missing from the Python API ### 3. Window Functions diff --git a/docs/source/user-guide/upgrade-guides.rst b/docs/source/user-guide/upgrade-guides.rst index 12d0f1bb3..9671594b8 100644 --- a/docs/source/user-guide/upgrade-guides.rst +++ b/docs/source/user-guide/upgrade-guides.rst @@ -45,6 +45,26 @@ After: config = SessionConfig().set("datafusion.execution.batch_size", "4096") ctx = SessionContext(config) +The aggregate functions :py:func:`~datafusion.functions.sum` and +:py:func:`~datafusion.functions.avg` now accept a ``distinct`` argument, matching +the other aggregate functions. ``distinct`` is inserted *before* ``filter`` in the +argument list, so any code that passed ``filter`` positionally must be updated to +pass it as a keyword argument. The types are distinct so a type checker should flag this. + +Before: + +.. code-block:: python + + f.sum(column("a"), my_filter) + f.avg(column("a"), my_filter) + +Now: + +.. code-block:: python + + f.sum(column("a"), filter=my_filter) + f.avg(column("a"), filter=my_filter) + DataFusion 53.0.0 ----------------- diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 383853007..c11a5c6cd 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -4563,6 +4563,7 @@ def grouping( def avg( expression: Expr, + distinct: bool = False, filter: Expr | None = None, ) -> Expr: """Returns the average value. @@ -4570,10 +4571,11 @@ def avg( This aggregate function expects a numeric expression and will return a float. If using the builder functions described in ref:`_aggregation` this function ignores - the options ``order_by``, ``null_treatment``, and ``distinct``. + the options ``order_by`` and ``null_treatment``. Args: expression: Values to combine into an array + distinct: If True, duplicate values are removed before averaging. filter: If provided, only compute against rows for which the filter is True Examples: @@ -4593,9 +4595,17 @@ def avg( ... ).alias("v")]) >>> result.collect_column("v")[0].as_py() 2.5 + + >>> df = ctx.from_pydict({"a": [1.0, 1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.avg( + ... dfn.col("a"), distinct=True, + ... ).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None - return Expr(f.avg(expression.expr, filter=filter_raw)) + return Expr(f.avg(expression.expr, distinct=distinct, filter=filter_raw)) def corr(value_y: Expr, value_x: Expr, filter: Expr | None = None) -> Expr: @@ -4880,6 +4890,7 @@ def min(expression: Expr, filter: Expr | None = None) -> Expr: def sum( expression: Expr, + distinct: bool = False, filter: Expr | None = None, ) -> Expr: """Computes the sum of a set of numbers. @@ -4887,10 +4898,11 @@ def sum( This aggregate function expects a numeric expression. If using the builder functions described in ref:`_aggregation` this function ignores - the options ``order_by``, ``null_treatment``, and ``distinct``. + the options ``order_by`` and ``null_treatment``. Args: expression: Values to combine into an array + distinct: If True, duplicate values are removed before summing. filter: If provided, only compute against rows for which the filter is True Examples: @@ -4910,9 +4922,17 @@ def sum( ... ).alias("v")]) >>> result.collect_column("v")[0].as_py() 5 + + >>> df = ctx.from_pydict({"a": [1, 1, 2, 3]}) + >>> result = df.aggregate( + ... [], [dfn.functions.sum( + ... dfn.col("a"), distinct=True, + ... ).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 6 """ filter_raw = filter.expr if filter is not None else None - return Expr(f.sum(expression.expr, filter=filter_raw)) + return Expr(f.sum(expression.expr, distinct=distinct, filter=filter_raw)) def stddev(expression: Expr, filter: Expr | None = None) -> Expr: diff --git a/python/tests/test_aggregation.py b/python/tests/test_aggregation.py index f5c54f756..ef51343aa 100644 --- a/python/tests/test_aggregation.py +++ b/python/tests/test_aggregation.py @@ -192,7 +192,9 @@ def test_aggregation_stats(df, agg_expr, calc_expected): False, ), (f.avg(column("b"), filter=column("a") != lit(1)), pa.array([5.0]), False), + (f.avg(column("b"), distinct=True), pa.array([5.0]), False), (f.sum(column("b"), filter=column("a") != lit(1)), pa.array([10]), False), + (f.sum(column("b"), distinct=True), pa.array([10]), False), (f.count(column("b"), distinct=True), pa.array([2]), False), (f.count(column("b"), filter=column("a") != 3), pa.array([2]), False), (f.count(), pa.array([3]), False), From 3d4c56c0757ddc0372ced03ffa97a13ea50d1bd8 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 29 May 2026 16:53:19 -0400 Subject: [PATCH 49/83] feat: create free-threaded python wheels (#1553) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Initial commit for free threaded python support * ci: use uvx to run maturin in native wheel builds The free-threaded matrix entries skip `uv sync` to avoid resolving project dependencies against cp313t/cp314t (many dev deps lack free-threaded wheels), so `uv run --no-project maturin` failed on macOS/Windows with "Failed to spawn: `maturin`". Switch to `uvx maturin@1.8.1`, which runs maturin in an isolated tool env independent of the project venv and matches the pin used by maturin-action for manylinux builds. Co-Authored-By: Claude Opus 4.7 (1M context) * ci: resolve free-threaded interpreter path explicitly on Windows maturin's `--interpreter python3.14t` fails on Windows because the free-threaded build ships as plain `python.exe` (no `tN` suffix). Look up `sys.executable` of the python on PATH (which actions/setup-python prepends with the free-threaded install), assert `Py_GIL_DISABLED == 1` so a misconfigured PATH can't silently build a GIL wheel, and normalize backslashes to forward slashes so the path survives re-expansion in the downstream `run:` line. Co-Authored-By: Claude Opus 4.7 (1M context) * build: enable PyO3 generate-import-lib for Windows free-threaded wheels Windows free-threaded Python does not expose `abiflags` in sysconfig, so PyO3's default Windows linkage path fails with "A python 3 interpreter on Windows does not define abiflags in its sysconfig ಠ_ಠ" when building cp31Xt wheels. Enabling the `generate-import-lib` PyO3 feature switches Windows builds to a generated import library (provided by the `python3-dll-a` crate) that does not depend on a fully populated sysconfig. It is a no-op on macOS and Linux and is compatible with the existing `abi3` feature. Co-Authored-By: Claude Opus 4.7 (1M context) * ci: bump maturin to 1.13.3 for Windows free-threaded support maturin 1.8.1 errors out on Windows free-threaded interpreters with "A python 3 interpreter on Windows does not define abiflags in its sysconfig" even when given a valid `python.exe`. Newer maturin releases handle the missing abiflags gracefully for cp31Xt builds. Bump both the `uvx maturin@` pin used for native macOS/Windows wheels and the `maturin-version` passed to PyO3/maturin-action for the manylinux containers. Co-Authored-By: Claude Opus 4.7 (1M context) * ci: standardize wheel build job names as " ()" The mac/Windows matrix shared a single name template that prepended "macOS arm64 & Windows" to every entry, which got truncated in the GitHub UI sidebar and made it hard to tell macOS and Windows runs apart. Rename all wheel build jobs to the same pattern so the OS, architecture, and python tag are visible at a glance: - Linux x86_64 / arm64 - macOS arm64 / x86_64 - Windows x86_64 Co-Authored-By: Claude Opus 4.7 (1M context) * taplo fmt * build: move pygithub to release group to fix free-threaded wheel builds pygithub pulls in cryptography via pyjwt[crypto]. cryptography 44.0.0 ships only abi3 wheels, which free-threaded interpreters cannot use, so uv builds it from sdist; its bundled PyO3 0.23.2 caps at Python 3.13 and fails on 3.14t. pygithub is only used by the manual release changelog script, so move it out of the dev group into a new release group. 'uv sync --dev' (used by CI test jobs) no longer drags in cryptography. * ci: pin uv venv to setup-python interpreter for free-threaded jobs Passing a bare version like '3.13t' to 'uv venv --python' let uv fall back to a different system interpreter (3.12), creating a venv whose ABI did not match the downloaded cp313t wheel and failing the install. Use the python-path output from setup-python so the venv uses exactly the interpreter that was set up. * taplo fmt * ci: set UV_PYTHON so uv sync keeps the free-threaded interpreter Pinning only 'uv venv --python' was not enough: 'uv sync' ignores the existing .venv, runs its own interpreter discovery, and recreated the venv with the system 3.12, again mismatching the cp313t wheel. Set UV_PYTHON to the setup-python interpreter for the install and test steps so every uv command (venv, sync, pip, run) uses it. * ci: run tests from the .venv, not the bare setup-python interpreter Setting UV_PYTHON on the test step pointed 'uv run --no-project pytest' at the setup-python interpreter, which has no pytest installed, causing 'Failed to spawn: pytest'. UV_PYTHON is only needed in the install step to build the .venv with the right interpreter; the test step must use that .venv. Drop UV_PYTHON from the test step. Co-Authored-By: Claude * ci: install datafusion wheel into the activated .venv Setting UV_PYTHON as a step env split the install across two environments: 'uv sync' populated .venv while 'uv pip install' targeted the bare setup-python interpreter, so the datafusion wheel never landed in .venv and 'import datafusion' failed under pytest. Pin the interpreter at 'uv venv --python', activate the venv, and pass --active to 'uv sync' so sync and pip install both target the same .venv. Co-Authored-By: Claude * ci: point uv at the venv interpreter by path for free-threaded jobs Activating the venv and passing --active still let 'uv sync' run its own interpreter discovery, which skips free-threaded builds and re-picked the system 3.12, recreating .venv and breaking the cp313t/cp314t wheel install. Pass the venv's own interpreter (.venv/bin/python) explicitly to 'uv sync', 'uv pip install', and 'uv run' so every step stays in the free-threaded environment created by 'uv venv'. Co-Authored-By: Claude --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .github/actions/build-wheel/action.yml | 110 +++++++++++++++++++++ .github/workflows/build.yml | 128 ++++++++++++++----------- .github/workflows/test.yml | 62 +++++++----- Cargo.lock | 10 ++ crates/core/Cargo.toml | 9 +- dev/release/README.md | 1 + pyproject.toml | 6 +- uv.lock | 6 +- 8 files changed, 247 insertions(+), 85 deletions(-) create mode 100644 .github/actions/build-wheel/action.yml diff --git a/.github/actions/build-wheel/action.yml b/.github/actions/build-wheel/action.yml new file mode 100644 index 000000000..25f75d1f8 --- /dev/null +++ b/.github/actions/build-wheel/action.yml @@ -0,0 +1,110 @@ +# 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. + +# Composite action that builds a datafusion-python wheel with maturin. +# Centralises the abi3-vs-free-threaded argument logic so platform jobs +# stay short and changes to wheel-build flags happen in one place. + +name: "Build wheel" +description: "Build datafusion-python wheel with maturin (abi3 or free-threaded)" + +inputs: + target: + description: "Rust target triple (e.g. x86_64-unknown-linux-gnu). Required when manylinux is set; ignored for native builds." + required: false + default: "" + python-tag: + description: "abi3 (covers 3.10..3.14 GIL builds) or a free-threaded interpreter such as 3.13t / 3.14t" + required: true + build-mode: + description: "release or debug" + required: true + features: + description: "Comma-separated extra features (in addition to those implied by the python-tag)" + required: false + default: "substrait" + manylinux: + description: "manylinux tag for maturin-action (e.g. 2_28). Leave empty to use uv-run maturin natively." + required: false + default: "" + out-dir: + description: "Output directory for built wheels" + required: false + default: "dist" + +outputs: + args: + description: "Computed maturin args (for debugging)" + value: ${{ steps.args.outputs.args }} + +runs: + using: "composite" + steps: + - name: Compute maturin args + id: args + shell: bash + run: | + set -euo pipefail + FEATURES="${{ inputs.features }}" + TAG="${{ inputs.python-tag }}" + if [ "$TAG" = "abi3" ]; then + # Default features include the `abi3` cargo feature. + # One wheel covers Python 3.10..3.14 (GIL builds only). + BUILD_ARGS="--features ${FEATURES}" + else + # Free-threaded build: disable abi3, force mimalloc back in, pin interpreter. + if [ "${RUNNER_OS:-}" = "Windows" ]; then + # Windows free-threaded builds ship as `python.exe` (no `tN` + # suffix). Resolve sys.executable so the path is independent of + # PATH ordering, and assert the interpreter is actually + # free-threaded before we hand the wheel off. + INTERP=$(python -c 'import sys; print(sys.executable)') + python -c "import sysconfig, sys; \ + v = sysconfig.get_config_var('Py_GIL_DISABLED'); \ + sys.exit(0 if v == 1 else f'expected free-threaded interpreter, got Py_GIL_DISABLED={v!r} at {sys.executable}')" + # Backslashes in BUILD_ARGS would be parsed as escapes when the + # output is re-expanded in the next step; use forward slashes + # (maturin/Rust accept them on Windows). + INTERP="${INTERP//\\//}" + else + INTERP="python${TAG}" + fi + BUILD_ARGS="--no-default-features --features mimalloc,${FEATURES} --interpreter ${INTERP}" + fi + if [ "${{ inputs.build-mode }}" = "release" ]; then + BUILD_ARGS="--release --strip ${BUILD_ARGS}" + fi + BUILD_ARGS="${BUILD_ARGS} --out ${{ inputs.out-dir }}" + echo "args=${BUILD_ARGS}" >> "$GITHUB_OUTPUT" + echo "maturin args: ${BUILD_ARGS}" + + - name: Build via maturin-action (manylinux container) + if: inputs.manylinux != '' + uses: PyO3/maturin-action@v1 + with: + target: ${{ inputs.target }} + manylinux: ${{ inputs.manylinux }} + maturin-version: "1.13.3" + args: ${{ steps.args.outputs.args }} + rustup-components: rust-std + + - name: Build via native maturin + if: inputs.manylinux == '' + shell: bash + # Use `uvx` so maturin is available even when `uv sync` was skipped + # (free-threaded matrix entries don't pre-populate the project venv). + run: uvx maturin@1.13.3 build ${{ steps.args.outputs.args }} diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 37a9dba03..593a343e1 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -135,8 +135,12 @@ jobs: # ============================================ build-manylinux-x86_64: needs: [generate-license, lint-rust, lint-python] - name: ManyLinux x86_64 + name: Linux x86_64 (${{ matrix.python-tag }}) runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-tag: ["abi3", "3.13t", "3.14t"] steps: - uses: actions/checkout@v6 @@ -153,7 +157,7 @@ jobs: - name: Cache Cargo uses: Swatinem/rust-cache@v2 with: - key: ${{ inputs.build_mode }} + key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: @@ -172,25 +176,18 @@ jobs: free -h swapon --show - - name: Build (release mode) - uses: PyO3/maturin-action@v1 - if: inputs.build_mode == 'release' - with: - target: x86_64-unknown-linux-gnu - manylinux: "2_28" - args: --release --strip --features protoc,substrait --out dist - rustup-components: rust-std - - - name: Build (debug mode) - uses: PyO3/maturin-action@v1 - if: inputs.build_mode == 'debug' + - name: Build wheel + uses: ./.github/actions/build-wheel with: target: x86_64-unknown-linux-gnu + python-tag: ${{ matrix.python-tag }} + build-mode: ${{ inputs.build_mode }} + features: "protoc,substrait" manylinux: "2_28" - args: --features protoc,substrait --out dist - rustup-components: rust-std + # FFI test wheel only needs to be built once per platform; gate to abi3. - name: Build FFI test library + if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: target: x86_64-unknown-linux-gnu @@ -202,10 +199,11 @@ jobs: - name: Archive wheels uses: actions/upload-artifact@v7 with: - name: dist-manylinux-x86_64 + name: dist-manylinux-x86_64-${{ matrix.python-tag }} path: dist/* - name: Archive FFI test wheel + if: matrix.python-tag == 'abi3' uses: actions/upload-artifact@v7 with: name: test-ffi-manylinux-x86_64 @@ -216,8 +214,12 @@ jobs: # ============================================ build-manylinux-aarch64: needs: [generate-license, lint-rust, lint-python] - name: ManyLinux arm64 + name: Linux arm64 (${{ matrix.python-tag }}) runs-on: ubuntu-24.04-arm + strategy: + fail-fast: false + matrix: + python-tag: ["abi3", "3.13t", "3.14t"] steps: - uses: actions/checkout@v6 @@ -234,7 +236,7 @@ jobs: - name: Cache Cargo uses: Swatinem/rust-cache@v2 with: - key: ${{ inputs.build_mode }} + key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: @@ -253,29 +255,20 @@ jobs: free -h swapon --show - - name: Build (release mode) - uses: PyO3/maturin-action@v1 - if: inputs.build_mode == 'release' - with: - target: aarch64-unknown-linux-gnu - manylinux: "2_28" - args: --release --strip --features protoc,substrait --out dist - rustup-components: rust-std - - - name: Build (debug mode) - uses: PyO3/maturin-action@v1 - if: inputs.build_mode == 'debug' + - name: Build wheel + uses: ./.github/actions/build-wheel with: target: aarch64-unknown-linux-gnu + python-tag: ${{ matrix.python-tag }} + build-mode: ${{ inputs.build_mode }} + features: "protoc,substrait" manylinux: "2_28" - args: --features protoc,substrait --out dist - rustup-components: rust-std - name: Archive wheels uses: actions/upload-artifact@v7 if: inputs.build_mode == 'release' with: - name: dist-manylinux-aarch64 + name: dist-manylinux-aarch64-${{ matrix.python-tag }} path: dist/* # ============================================ @@ -283,13 +276,13 @@ jobs: # ============================================ build-python-mac-win: needs: [generate-license, lint-rust, lint-python] - name: macOS arm64 & Windows + name: ${{ matrix.os == 'macos-latest' && 'macOS arm64' || 'Windows x86_64' }} (${{ matrix.python-tag }}) runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: - python-version: ["3.10"] os: [macos-latest, windows-latest] + python-tag: ["abi3", "3.13t", "3.14t"] steps: - uses: actions/checkout@v6 @@ -305,7 +298,14 @@ jobs: - name: Cache Cargo uses: Swatinem/rust-cache@v2 with: - key: ${{ inputs.build_mode }} + key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} + + - name: Setup Python (free-threaded) + if: matrix.python-tag != 'abi3' + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-tag }} + freethreaded: true - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: @@ -318,22 +318,22 @@ jobs: repo-token: ${{ secrets.GITHUB_TOKEN }} - name: Install dependencies + if: matrix.python-tag == 'abi3' run: uv sync --dev --no-install-package datafusion - # Run clippy BEFORE maturin so we can avoid rebuilding. The features must match - # exactly the features used by maturin. Linux maturin builds need to happen in a - # container so only run this for our mac runner. + # Clippy is interpreter-agnostic; run once per OS (against the abi3 entry) + # so the matrix doesn't pay the cost three times. - name: Run Clippy - if: matrix.os != 'windows-latest' + if: matrix.os != 'windows-latest' && matrix.python-tag == 'abi3' run: cargo clippy --no-deps --all-targets --features substrait -- -D warnings - - name: Build Python package (release mode) - if: inputs.build_mode == 'release' - run: uv run --no-project maturin build --release --strip --features substrait - - - name: Build Python package (debug mode) - if: inputs.build_mode != 'release' - run: uv run --no-project maturin build --features substrait + - name: Build wheel + uses: ./.github/actions/build-wheel + with: + python-tag: ${{ matrix.python-tag }} + build-mode: ${{ inputs.build_mode }} + features: "substrait" + out-dir: "target/wheels" - name: List Windows wheels if: matrix.os == 'windows-latest' @@ -350,7 +350,7 @@ jobs: uses: actions/upload-artifact@v7 if: inputs.build_mode == 'release' with: - name: dist-${{ matrix.os }} + name: dist-${{ matrix.os }}-${{ matrix.python-tag }} path: target/wheels/* # ============================================ @@ -359,11 +359,12 @@ jobs: build-macos-x86_64: if: inputs.build_mode == 'release' needs: [generate-license, lint-rust, lint-python] + name: macOS x86_64 (${{ matrix.python-tag }}) runs-on: macos-15-intel strategy: fail-fast: false matrix: - python-version: ["3.10"] + python-tag: ["abi3", "3.13t", "3.14t"] steps: - uses: actions/checkout@v6 @@ -379,7 +380,14 @@ jobs: - name: Cache Cargo uses: Swatinem/rust-cache@v2 with: - key: ${{ inputs.build_mode }} + key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} + + - name: Setup Python (free-threaded) + if: matrix.python-tag != 'abi3' + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-tag }} + freethreaded: true - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: @@ -392,11 +400,16 @@ jobs: repo-token: ${{ secrets.GITHUB_TOKEN }} - name: Install dependencies + if: matrix.python-tag == 'abi3' run: uv sync --dev --no-install-package datafusion - - name: Build (release mode) - run: | - uv run --no-project maturin build --release --strip --features substrait + - name: Build wheel + uses: ./.github/actions/build-wheel + with: + python-tag: ${{ matrix.python-tag }} + build-mode: ${{ inputs.build_mode }} + features: "substrait" + out-dir: "target/wheels" - name: List Mac wheels run: find target/wheels/ @@ -404,7 +417,7 @@ jobs: - name: Archive wheels uses: actions/upload-artifact@v7 with: - name: dist-macos-aarch64 + name: dist-macos-aarch64-${{ matrix.python-tag }} path: target/wheels/* # ============================================ @@ -509,11 +522,12 @@ jobs: with: enable-cache: true - # Download the Linux wheel built in the previous job + # Download the Linux wheel built in the previous job. + # Docs only need the abi3 wheel — interpreter doesn't matter for sphinx. - name: Download pre-built Linux wheel uses: actions/download-artifact@v8 with: - name: dist-manylinux-x86_64 + name: dist-manylinux-x86_64-abi3 path: wheels/ # Install from the pre-built wheels diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2cd792ea9..0c8fa4f79 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -15,8 +15,9 @@ # specific language governing permissions and limitations # under the License. -# Reusable workflow for running tests -# This ensures the same tests run for both debug (PRs) and release (main/tags) builds +# Reusable workflow for running tests. +# Single matrix covers both GIL (abi3 wheel) and free-threaded +# (per-interpreter wheels) builds. name: Test @@ -35,61 +36,69 @@ jobs: strategy: fail-fast: false matrix: - python-version: - - "3.10" - - "3.11" - - "3.12" - - "3.13" - - "3.14" - toolchain: - - "stable" - + include: + # GIL builds — all share the same abi3 wheel. + - { python-version: "3.10", wheel-tag: "abi3", freethreaded: false } + - { python-version: "3.11", wheel-tag: "abi3", freethreaded: false } + - { python-version: "3.12", wheel-tag: "abi3", freethreaded: false } + - { python-version: "3.13", wheel-tag: "abi3", freethreaded: false } + - { python-version: "3.14", wheel-tag: "abi3", freethreaded: false } + # Free-threaded builds — one wheel per interpreter. + - { python-version: "3.13t", wheel-tag: "3.13t", freethreaded: true } + - { python-version: "3.14t", wheel-tag: "3.14t", freethreaded: true } steps: - uses: actions/checkout@v6 - name: Setup Python + id: setup-python uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} + freethreaded: ${{ matrix.freethreaded }} - name: Cache Cargo uses: actions/cache@v5 with: path: ~/.cargo - key: cargo-cache-${{ matrix.toolchain }}-${{ hashFiles('Cargo.lock') }} + key: cargo-cache-stable-${{ hashFiles('Cargo.lock') }} - name: Install dependencies uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b with: enable-cache: true - # Download the Linux wheel built in the build workflow - name: Download pre-built Linux wheel uses: actions/download-artifact@v8 with: - name: dist-manylinux-x86_64 + name: dist-manylinux-x86_64-${{ matrix.wheel-tag }} path: wheels/ - # Download the FFI test wheel + # FFI test wheel only built once (under the abi3 matrix entry in build.yml). - name: Download pre-built FFI test wheel + if: matrix.wheel-tag == 'abi3' uses: actions/download-artifact@v8 with: name: test-ffi-manylinux-x86_64 path: wheels/ - # Install from the pre-built wheels - name: Install from pre-built wheels run: | set -x - uv venv - # Install development dependencies - uv sync --dev --no-install-package datafusion - # Install all pre-built wheels + # Create the venv with the setup-python interpreter, then point + # every uv command explicitly at the venv's own interpreter. + # uv's interpreter discovery skips free-threaded builds unless + # asked by exact path, so plain `uv sync` (even with an activated + # venv or --active) re-picks the system 3.12 and recreates .venv. + # Targeting .venv/bin/python keeps sync and pip install in the + # same 3.13t/3.14t environment as the cp313t/cp314t wheel. + uv venv --python "${{ steps.setup-python.outputs.python-path }}" + VENV_PY="$PWD/.venv/bin/python" + uv sync --python "$VENV_PY" --dev --no-install-package datafusion WHEELS=$(find wheels/ -name "*.whl") if [ -n "$WHEELS" ]; then echo "Installing wheels:" echo "$WHEELS" - uv pip install wheels/*.whl + uv pip install --python "$VENV_PY" wheels/*.whl else echo "ERROR: No wheels found!" exit 1 @@ -98,16 +107,24 @@ jobs: - name: Run tests env: RUST_BACKTRACE: 1 + # On free-threaded interpreters, fail loud if any C extension + # re-enables the GIL implicitly. + PYTHON_GIL: ${{ matrix.freethreaded && '0' || '' }} run: | git submodule update --init - uv run --no-project pytest -v --import-mode=importlib + # Use the .venv interpreter directly; uv discovery would skip the + # free-threaded build and re-pick the system 3.12 (see install step). + uv run --python "$PWD/.venv/bin/python" --no-project pytest -v --import-mode=importlib + # FFI + TPC-H examples only need to run once; gate to abi3 entries. - name: FFI unit tests + if: matrix.wheel-tag == 'abi3' run: | cd examples/datafusion-ffi-example uv run --no-project pytest python/tests/_test*.py - name: Run tpchgen-cli to create 1 Gb dataset + if: matrix.wheel-tag == 'abi3' run: | mkdir examples/tpch/data cd examples/tpch/data @@ -115,6 +132,7 @@ jobs: uv run --no-project tpchgen-cli -s 1 --format=parquet - name: Run TPC-H examples + if: matrix.wheel-tag == 'abi3' run: | cd examples/tpch uv run --no-project pytest _tests.py diff --git a/Cargo.lock b/Cargo.lock index 6a1ef2447..d3cedb628 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2904,6 +2904,7 @@ version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" dependencies = [ + "python3-dll-a", "target-lexicon", ] @@ -2953,6 +2954,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "python3-dll-a" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d80ba7540edb18890d444c5aa8e1f1f99b1bdf26fb26ae383135325f4a36042b" +dependencies = [ + "cc", +] + [[package]] name = "quick-xml" version = "0.39.4" diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index d714dc978..1f5b4e305 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -42,8 +42,7 @@ tokio = { workspace = true, features = [ ] } pyo3 = { workspace = true, features = [ "extension-module", - "abi3", - "abi3-py310", + "generate-import-lib", ] } pyo3-async-runtimes = { workspace = true, features = ["tokio-runtime"] } pyo3-log = { workspace = true } @@ -74,7 +73,11 @@ prost-types = { workspace = true } pyo3-build-config = { workspace = true } [features] -default = ["mimalloc"] +default = ["mimalloc", "abi3"] +# Stable ABI build — single wheel covers Python 3.10..3.14 (GIL builds only). +# Mutually exclusive with free-threaded interpreters (cp313t / cp314t); the +# free-threaded wheel build must pass --no-default-features. +abi3 = ["pyo3/abi3", "pyo3/abi3-py310"] protoc = ["datafusion-substrait/protoc"] substrait = ["dep:datafusion-substrait"] diff --git a/dev/release/README.md b/dev/release/README.md index 1b02ca576..384ef9210 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -75,6 +75,7 @@ We maintain a `CHANGELOG.md` so our users know what has been changed between rel The changelog is generated using a Python script: ```bash +$ uv sync --group release $ GITHUB_TOKEN= ./dev/release/generate-changelog.py 52.0.0 HEAD 53.0.0 > dev/changelog/53.0.0.md ``` diff --git a/pyproject.toml b/pyproject.toml index 418640a49..2b6a976db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ classifiers = [ "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: Free Threading :: 2 - Beta", "Programming Language :: Python", "Programming Language :: Rust", ] @@ -201,7 +202,6 @@ dev = [ "numpy>=2.3.2;python_version>='3.14'", "pre-commit>=4.3.0", "pyarrow>=19.0.0", - "pygithub==2.5.0", "pytest-asyncio>=0.23.3", "pytest-timeout>=2.3.1", "pytest>=7.4.4", @@ -209,6 +209,10 @@ dev = [ "ruff>=0.15.1", "toml>=0.10.2", ] +# Release tooling only. Kept out of `dev` because pygithub pulls in +# cryptography, which ships no free-threaded wheel and fails to build +# from sdist under free-threaded interpreters (PyO3 < 3.14 support). +release = ["pygithub==2.5.0"] docs = [ "ipython>=8.12.3", "jinja2>=3.1.5", diff --git a/uv.lock b/uv.lock index 26ab8b20e..6673b7fe2 100644 --- a/uv.lock +++ b/uv.lock @@ -340,7 +340,6 @@ dev = [ { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, { name = "pre-commit" }, { name = "pyarrow" }, - { name = "pygithub" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-timeout" }, @@ -359,6 +358,9 @@ docs = [ { name = "sphinx" }, { name = "sphinx-autoapi" }, ] +release = [ + { name = "pygithub" }, +] [package.metadata] requires-dist = [ @@ -378,7 +380,6 @@ dev = [ { name = "numpy", marker = "python_full_version >= '3.14'", specifier = ">=2.3.2" }, { name = "pre-commit", specifier = ">=4.3.0" }, { name = "pyarrow", specifier = ">=19.0.0" }, - { name = "pygithub", specifier = "==2.5.0" }, { name = "pytest", specifier = ">=7.4.4" }, { name = "pytest-asyncio", specifier = ">=0.23.3" }, { name = "pytest-timeout", specifier = ">=2.3.1" }, @@ -397,6 +398,7 @@ docs = [ { name = "sphinx", specifier = ">=7.1.2" }, { name = "sphinx-autoapi", specifier = ">=3.4.0" }, ] +release = [{ name = "pygithub", specifier = "==2.5.0" }] [[package]] name = "decorator" From af388667f8d8583abdd634fa343252ddb4a0da16 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 29 May 2026 18:04:24 -0400 Subject: [PATCH 50/83] feat: expose lambda and higher-order array functions (#1561) * feat: expose lambda and higher-order array functions Add a Pythonic API for DataFusion's higher-order array functions and the lambda expressions they consume. - Rust: lambda_, lambda_var, array_transform, and array_any_match pyfunctions, plus a ResolveLambdaVariables analyzer rule so expression-builder plans (which emit unresolved lambda variables) resolve before optimization. - Python: array_transform / array_any_match (with list_transform, any_match, list_any_match aliases) accept either a Python callable or an explicit lambda built with lambda_ / lambda_var. Callables are introspected so their parameter names become the lambda parameters. - Tests and docs (expressions guide + agent skill), noting v1 limits: lambda expressions are not serializable, and SQL arrow syntax needs the DuckDB dialect. * test: fold lambda tests into pytest parameterization Combine the eight higher-order function result tests into a single parametrized test_higher_order_function_results, and the two to_lambda rejection tests into test_to_lambda_rejects_invalid_arg. Each case keeps a readable id via pytest.param. Co-Authored-By: Claude * feat: expose array_filter higher-order function Add array_filter, the remaining lambda-based higher-order array function in DataFusion (alongside the already-exposed array_transform and array_any_match). Includes the list_filter alias matching upstream, tests, and documentation in the expressions guide and skill. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: emphasize lambda terminology, trim skill lambda section Lead user-facing array-lambda docs with "lambda function" instead of "higher-order function," which is less recognizable to users. Drop the alias list, serialization caveat, and DuckDB-dialect note from the skill to keep it lean; those details already live in the docstrings. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: broaden SQL lambda dialect coverage Other dialects (ClickHouse, Snowflake, Databricks) also enable lambda parsing via sqlparser-rs. Document the full set and recommend the ``lambda x: x`` keyword form, since DuckDB will drop the ``x -> x`` arrow form in v2.1. Parametrize the SQL test over the four dialects using the keyword syntax. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude --- crates/core/src/analyzer.rs | 59 +++++ crates/core/src/context.rs | 1 + crates/core/src/functions.rs | 45 ++++ crates/core/src/lib.rs | 1 + .../common-operations/expressions.rst | 49 ++++ python/datafusion/functions.py | 218 +++++++++++++++++- python/tests/test_lambda.py | 144 ++++++++++++ skills/datafusion_python/SKILL.md | 18 ++ 8 files changed, 534 insertions(+), 1 deletion(-) create mode 100644 crates/core/src/analyzer.rs create mode 100644 python/tests/test_lambda.py diff --git a/crates/core/src/analyzer.rs b/crates/core/src/analyzer.rs new file mode 100644 index 000000000..3a77e08a3 --- /dev/null +++ b/crates/core/src/analyzer.rs @@ -0,0 +1,59 @@ +// 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. + +//! Analyzer rules layered on top of DataFusion's defaults. + +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::logical_expr::LogicalPlan; +use datafusion::optimizer::AnalyzerRule; + +/// Resolve [`LambdaVariable`] references into bound lambda parameters. +/// +/// DataFusion's SQL planner resolves lambda variables inline as it plans a +/// higher-order function call, so SQL-built plans never carry unresolved +/// variables. Plans assembled programmatically through the Python expression +/// builder (e.g. `array_transform(col("xs"), lambda_(["v"], lambda_var("v")))`) +/// do carry them, and nothing in the default analyzer resolves them. This rule +/// runs [`LogicalPlan::resolve_lambda_variables`] so both construction paths +/// reach the optimizer with bound lambdas. +/// +/// [`LambdaVariable`]: datafusion::logical_expr::expr::LambdaVariable +#[derive(Debug)] +pub struct ResolveLambdaVariables {} + +impl ResolveLambdaVariables { + pub fn new() -> Self { + Self {} + } +} + +impl Default for ResolveLambdaVariables { + fn default() -> Self { + Self::new() + } +} + +impl AnalyzerRule for ResolveLambdaVariables { + fn analyze(&self, plan: LogicalPlan, _config: &ConfigOptions) -> Result { + plan.resolve_lambda_variables().map(|t| t.data) + } + + fn name(&self) -> &str { + "resolve_lambda_variables" + } +} diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 4606246cf..d714861a6 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -396,6 +396,7 @@ impl PySessionContext { .with_config(config) .with_runtime_env(runtime) .with_default_features() + .with_analyzer_rule(Arc::new(crate::analyzer::ResolveLambdaVariables::new())) .build(); let ctx = Arc::new(SessionContext::new_with_state(session_state)); Ok(PySessionContext { diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 5f47d123b..395d5ebfd 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -159,6 +159,44 @@ fn array_slice(array: PyExpr, begin: PyExpr, end: PyExpr, stride: Option .into() } +/// Create a lambda expression from a list of parameter names and a body +/// expression. The body should reference the parameters via [`lambda_var`]. +/// Exposed to Python as `lambda_` because `lambda` is a reserved keyword. +#[pyfunction] +#[pyo3(name = "lambda_")] +fn py_lambda(params: Vec, body: PyExpr) -> PyExpr { + datafusion::logical_expr::lambda(params, body.into()).into() +} + +/// Create an unresolved lambda variable reference by name. The owning +/// higher-order function resolves it against its lambda parameters during +/// planning. +#[pyfunction] +fn lambda_var(name: String) -> PyExpr { + datafusion::logical_expr::lambda_var(name).into() +} + +/// Higher-order function: apply `transform` (a lambda) to each element of +/// `array`, returning a new array of the results. +#[pyfunction] +fn array_transform(array: PyExpr, transform: PyExpr) -> PyExpr { + datafusion::functions_nested::expr_fn::array_transform(array.into(), transform.into()).into() +} + +/// Higher-order function: return true if any element of `array` satisfies +/// `predicate` (a lambda returning a boolean). +#[pyfunction] +fn array_any_match(array: PyExpr, predicate: PyExpr) -> PyExpr { + datafusion::functions_nested::expr_fn::array_any_match(array.into(), predicate.into()).into() +} + +/// Higher-order function: keep the elements of `array` for which `predicate` +/// (a lambda returning a boolean) is true, returning a new filtered array. +#[pyfunction] +fn array_filter(array: PyExpr, predicate: PyExpr) -> PyExpr { + datafusion::functions_nested::expr_fn::array_filter(array.into(), predicate.into()).into() +} + /// Computes a binary hash of the given data. type is the algorithm to use. /// Standard algorithms are md5, sha224, sha256, sha384, sha512, blake2s, blake2b, and blake3. // #[pyfunction(value, method)] @@ -1082,6 +1120,13 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(encode))?; m.add_wrapped(wrap_pyfunction!(decode))?; + // Lambda / higher-order functions + m.add_wrapped(wrap_pyfunction!(py_lambda))?; + m.add_wrapped(wrap_pyfunction!(lambda_var))?; + m.add_wrapped(wrap_pyfunction!(array_transform))?; + m.add_wrapped(wrap_pyfunction!(array_any_match))?; + m.add_wrapped(wrap_pyfunction!(array_filter))?; + // Array Functions m.add_wrapped(wrap_pyfunction!(array_append))?; m.add_wrapped(wrap_pyfunction!(array_concat))?; diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 79bf77717..48abcedc9 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -27,6 +27,7 @@ use mimalloc::MiMalloc; use pyo3::prelude::*; #[allow(clippy::borrow_deref_ref)] +pub mod analyzer; pub mod catalog; pub mod codec; pub mod common; diff --git a/docs/source/user-guide/common-operations/expressions.rst b/docs/source/user-guide/common-operations/expressions.rst index ae1ccc0dc..f52c79ddb 100644 --- a/docs/source/user-guide/common-operations/expressions.rst +++ b/docs/source/user-guide/common-operations/expressions.rst @@ -145,6 +145,55 @@ 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]]`. +Lambda functions +---------------- + +Some array functions take a *lambda function*: a small function that runs once +per element. :py:func:`~datafusion.functions.array_transform` maps a lambda over +every element, :py:func:`~datafusion.functions.array_filter` keeps the elements +for which a predicate lambda is true, and +:py:func:`~datafusion.functions.array_any_match` returns whether any element +satisfies a predicate lambda. (Functions that take another function as an +argument are sometimes called *higher-order* functions.) + +The simplest way to supply a lambda is a Python ``lambda``. Its parameter names +become the lambda parameters, and its return value becomes the body. + +.. ipython:: python + + from datafusion import SessionContext, col + from datafusion import functions as f + + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1, 2, 3], [4, 5]]}) + df.select(f.array_transform(col("a"), lambda v: v * 2).alias("doubled")) + df.select(f.array_filter(col("a"), lambda v: v > 2).alias("big_only")) + df.select(f.array_any_match(col("a"), lambda v: v > 3).alias("has_big")) + +If you need explicit control over parameter names, build the lambda with +:py:func:`~datafusion.functions.lambda_` and reference its parameters with +:py:func:`~datafusion.functions.lambda_var`. The following is equivalent to the +``array_transform`` call above. + +.. ipython:: python + + from datafusion import lit + + double_fn = f.lambda_(["v"], f.lambda_var("v") * lit(2)) + df.select(f.array_transform(col("a"), double_fn).alias("doubled")) + +.. note:: + + Lambda expressions cannot yet be serialized: calling + :py:meth:`~datafusion.expr.Expr.to_bytes` or pickling an expression that + contains a lambda raises ``Lambda not implemented``. SQL lambda syntax is + only parsed by dialects that support lambdas; set + ``datafusion.sql_parser.dialect`` to one of ``DuckDB``, ``ClickHouse``, + ``Snowflake``, or ``Databricks``. Both arrow syntax (``x -> x * 2``) and + keyword syntax (``lambda x: x * 2``) parse. DuckDB will drop the arrow + form in v2.1, so prefer ``lambda x: x * 2`` for forward compatibility. + The Python expression builder shown above works regardless of dialect. + Testing membership in a list ---------------------------- diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index c11a5c6cd..c8f07497d 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -38,10 +38,14 @@ from __future__ import annotations -from typing import Any +import inspect +from typing import TYPE_CHECKING, Any import pyarrow as pa +if TYPE_CHECKING: + from collections.abc import Callable + from datafusion._internal import functions as f from datafusion.common import NullTreatment from datafusion.expr import ( @@ -61,12 +65,14 @@ "acos", "acosh", "alias", + "any_match", "approx_distinct", "approx_median", "approx_percentile_cont", "approx_percentile_cont_with_weight", "array", "array_agg", + "array_any_match", "array_any_value", "array_append", "array_cat", @@ -79,6 +85,7 @@ "array_empty", "array_except", "array_extract", + "array_filter", "array_has", "array_has_all", "array_has_any", @@ -108,6 +115,7 @@ "array_slice", "array_sort", "array_to_string", + "array_transform", "array_union", "arrays_overlap", "arrays_zip", @@ -188,6 +196,8 @@ "isnan", "iszero", "lag", + "lambda_", + "lambda_var", "last_value", "lcm", "lead", @@ -195,6 +205,7 @@ "left", "length", "levenshtein", + "list_any_match", "list_any_value", "list_append", "list_cat", @@ -207,6 +218,7 @@ "list_empty", "list_except", "list_extract", + "list_filter", "list_has", "list_has_all", "list_has_any", @@ -237,6 +249,7 @@ "list_slice", "list_sort", "list_to_string", + "list_transform", "list_union", "list_zip", "ln", @@ -459,6 +472,209 @@ def list_join(expr: Expr, delimiter: Expr | str) -> Expr: return array_to_string(expr, delimiter) +def lambda_var(name: str) -> Expr: + """Create an unresolved reference to a lambda parameter by ``name``. + + Use this inside the body passed to :py:func:`lambda_` to refer to one of the + lambda's parameters. The owning higher-order function (such as + :py:func:`array_transform`) binds the variable to a concrete element type + during query planning. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> double_fn = F.lambda_(["v"], F.lambda_var("v") * lit(2)) + >>> df.select( + ... F.array_transform(col("a"), double_fn).alias("d") + ... ).collect_column("d")[0].as_py() + [2, 4, 6] + + See Also: + :py:func:`lambda_`, :py:func:`array_transform`, :py:func:`array_any_match`. + """ + return Expr(f.lambda_var(name)) + + +def lambda_(params: list[str], body: Expr) -> Expr: + """Create a lambda expression from parameter names and a body expression. + + This is the explicit form of building a lambda. Most callers can instead + pass a Python callable directly to a higher-order function such as + :py:func:`array_transform`, which builds the lambda automatically. Reach for + ``lambda_`` when you want explicit control over the parameter names. + + Args: + params: Ordered lambda parameter names. + body: Body expression that references the parameters via + :py:func:`lambda_var`. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> double_fn = F.lambda_(["v"], F.lambda_var("v") * lit(2)) + >>> df.select( + ... F.array_transform(col("a"), double_fn).alias("d") + ... ).collect_column("d")[0].as_py() + [2, 4, 6] + + See Also: + :py:func:`lambda_var`, :py:func:`array_transform`, :py:func:`array_any_match`. + """ + return Expr(f.lambda_(params, body.expr)) + + +def _to_lambda(fn: Expr | Callable[..., Any]) -> Expr: + """Coerce ``fn`` to a lambda ``Expr``. + + Accepts either an ``Expr`` produced by :py:func:`lambda_` (returned + unchanged) or a Python callable. A callable is introspected for its + parameter names; those names become :py:func:`lambda_var` references passed + positionally into the callable, and its return value (coerced to an + ``Expr``) becomes the lambda body. + """ + if isinstance(fn, Expr): + return fn + if not callable(fn): + msg = f"expected an Expr or callable, got {type(fn).__name__}" + raise TypeError(msg) + params = list(inspect.signature(fn).parameters) + if not params: + msg = "lambda callable must accept at least one parameter" + raise ValueError(msg) + body = coerce_to_expr(fn(*[lambda_var(p) for p in params])) + return lambda_(params, body) + + +def array_transform(array: Expr, transform: Expr | Callable[..., Any]) -> Expr: + """Transform each element of ``array`` with a lambda. + + ``transform`` may be a Python callable, which is converted to a lambda + automatically (its parameter names become the lambda parameters), or an + explicit lambda built with :py:func:`lambda_`. + + Examples: + Using a Python callable: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> df.select( + ... F.array_transform(col("a"), lambda v: v * 2).alias("d") + ... ).collect_column("d")[0].as_py() + [2, 4, 6] + + Using an explicit lambda built with :py:func:`lambda_`: + + >>> double_fn = F.lambda_(["v"], F.lambda_var("v") * lit(2)) + >>> df.select( + ... F.array_transform(col("a"), double_fn).alias("d") + ... ).collect_column("d")[0].as_py() + [2, 4, 6] + + See Also: + :py:func:`array_any_match`, :py:func:`lambda_`. + """ + return Expr(f.array_transform(array.expr, _to_lambda(transform).expr)) + + +def list_transform(array: Expr, transform: Expr | Callable[..., Any]) -> Expr: + """Transform each element of a list with a lambda. + + See Also: + This is an alias for :py:func:`array_transform`. + """ + return array_transform(array, transform) + + +def array_any_match(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr: + """Return ``True`` if any element of ``array`` satisfies ``predicate``. + + ``predicate`` may be a Python callable, converted to a lambda + automatically, or an explicit lambda built with :py:func:`lambda_`. It must + return a boolean expression. + + Examples: + Using a Python callable: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> df.select( + ... F.array_any_match(col("a"), lambda v: v > 2).alias("m") + ... ).collect_column("m")[0].as_py() + True + + Using an explicit lambda built with :py:func:`lambda_`: + + >>> predicate = F.lambda_(["v"], F.lambda_var("v") > lit(2)) + >>> df.select( + ... F.array_any_match(col("a"), predicate).alias("m") + ... ).collect_column("m")[0].as_py() + True + + See Also: + :py:func:`array_transform`, :py:func:`lambda_`. + """ + return Expr(f.array_any_match(array.expr, _to_lambda(predicate).expr)) + + +def any_match(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr: + """Return ``True`` if any element of an array satisfies a predicate. + + See Also: + This is an alias for :py:func:`array_any_match`. + """ + return array_any_match(array, predicate) + + +def list_any_match(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr: + """Return ``True`` if any element of a list satisfies a predicate. + + See Also: + This is an alias for :py:func:`array_any_match`. + """ + return array_any_match(array, predicate) + + +def array_filter(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr: + """Keep the elements of ``array`` for which ``predicate`` is ``True``. + + ``predicate`` may be a Python callable, converted to a lambda + automatically, or an explicit lambda built with :py:func:`lambda_`. It must + return a boolean expression. The result is a new array containing only the + matching elements. + + Examples: + Using a Python callable: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3, 4, 5]]}) + >>> df.select( + ... F.array_filter(col("a"), lambda v: v > 2).alias("f") + ... ).collect_column("f")[0].as_py() + [3, 4, 5] + + Using an explicit lambda built with :py:func:`lambda_`: + + >>> predicate = F.lambda_(["v"], F.lambda_var("v") > lit(2)) + >>> df.select( + ... F.array_filter(col("a"), predicate).alias("f") + ... ).collect_column("f")[0].as_py() + [3, 4, 5] + + See Also: + :py:func:`array_transform`, :py:func:`array_any_match`, :py:func:`lambda_`. + """ + return Expr(f.array_filter(array.expr, _to_lambda(predicate).expr)) + + +def list_filter(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr: + """Keep the elements of a list for which a predicate is ``True``. + + See Also: + This is an alias for :py:func:`array_filter`. + """ + return array_filter(array, predicate) + + def in_list(arg: Expr, values: list[Expr], negated: bool = False) -> Expr: """Returns whether the argument is contained within the list ``values``. diff --git a/python/tests/test_lambda.py b/python/tests/test_lambda.py new file mode 100644 index 000000000..68be22a04 --- /dev/null +++ b/python/tests/test_lambda.py @@ -0,0 +1,144 @@ +# 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. +"""Tests for lambda expressions and higher-order array functions.""" + +import pytest +from datafusion import SessionConfig, SessionContext, col, lit +from datafusion import functions as f + + +@pytest.fixture +def df(): + ctx = SessionContext() + return ctx.from_pydict({"a": [[1, 2, 3], [4, 5]]}) + + +def _column(df, expr, name): + return df.select(expr.alias(name)).collect_column(name).to_pylist() + + +@pytest.mark.parametrize( + ("build_expr", "expected"), + [ + pytest.param( + lambda: f.array_transform(col("a"), lambda v: v * 2), + [[2, 4, 6], [8, 10]], + id="array_transform_callable", + ), + pytest.param( + lambda: f.array_transform( + col("a"), f.lambda_(["v"], f.lambda_var("v") * lit(2)) + ), + [[2, 4, 6], [8, 10]], + id="array_transform_explicit_lambda", + ), + pytest.param( + lambda: f.array_transform(col("a"), lambda v: 0), + [[0, 0, 0], [0, 0]], + id="array_transform_literal_body_is_coerced", + ), + pytest.param( + lambda: f.list_transform(col("a"), lambda v: v + 1), + [[2, 3, 4], [5, 6]], + id="list_transform_alias", + ), + pytest.param( + lambda: f.array_any_match(col("a"), lambda v: v > 3), + [False, True], + id="array_any_match_callable", + ), + pytest.param( + lambda: f.array_any_match( + col("a"), f.lambda_(["v"], f.lambda_var("v") > lit(2)) + ), + [True, True], + id="array_any_match_explicit_lambda", + ), + pytest.param( + lambda: f.any_match(col("a"), lambda v: v > 4), + [False, True], + id="any_match_alias", + ), + pytest.param( + lambda: f.list_any_match(col("a"), lambda v: v > 4), + [False, True], + id="list_any_match_alias", + ), + pytest.param( + lambda: f.array_filter(col("a"), lambda v: v > 2), + [[3], [4, 5]], + id="array_filter_callable", + ), + pytest.param( + lambda: f.array_filter( + col("a"), f.lambda_(["v"], f.lambda_var("v") > lit(2)) + ), + [[3], [4, 5]], + id="array_filter_explicit_lambda", + ), + pytest.param( + lambda: f.list_filter(col("a"), lambda v: v > 2), + [[3], [4, 5]], + id="list_filter_alias", + ), + ], +) +def test_higher_order_function_results(df, build_expr, expected): + assert _column(df, build_expr(), "r") == expected + + +def test_lambda_param_name_appears_in_plan(df): + # The user-chosen parameter name should survive into the displayed plan + # rather than a synthetic placeholder. + expr = f.array_transform(col("a"), lambda value: value * 2) + assert "value" in expr.canonical_name() + + +@pytest.mark.parametrize( + ("arg", "exc_type", "match"), + [ + pytest.param(42, TypeError, "expected an Expr or callable", id="non_callable"), + pytest.param( + lambda: lit(1), + ValueError, + "at least one parameter", + id="zero_arg_callable", + ), + ], +) +def test_to_lambda_rejects_invalid_arg(arg, exc_type, match): + with pytest.raises(exc_type, match=match): + f.array_transform(col("a"), arg) + + +@pytest.mark.parametrize("dialect", ["DuckDB", "ClickHouse", "Snowflake", "Databricks"]) +def test_sql_lambda_keyword_syntax(dialect): + # ``lambda x: x * 2`` is the forward-compatible syntax. DuckDB will drop + # the arrow form (``x -> ...``) in v2.1; the keyword form is supported by + # every dialect in sqlparser-rs that enables lambda functions. + ctx = SessionContext(SessionConfig().set("datafusion.sql_parser.dialect", dialect)) + result = ctx.sql( + "select array_transform([1, 2, 3], lambda x: x * 2) as d" + ).collect_column("d") + assert result.to_pylist() == [[2, 4, 6]] + + +def test_pickle_lambda_expr_not_supported(): + # v1 limitation: upstream proto serialization rejects lambda expressions. + expr = f.array_transform(col("a"), lambda v: v * 2) + with pytest.raises(Exception, match="Lambda not implemented"): + expr.to_bytes() diff --git a/skills/datafusion_python/SKILL.md b/skills/datafusion_python/SKILL.md index 98fa2c7aa..1aeb78777 100644 --- a/skills/datafusion_python/SKILL.md +++ b/skills/datafusion_python/SKILL.md @@ -488,6 +488,24 @@ col("array_col")[0] # access array element (0-indexed) col("array_col")[1:3] # array slice (0-indexed) ``` +### Lambda Functions + +Some array functions take a lambda function that runs once per element. Pass a +Python `lambda` directly — its parameter names become the lambda parameters and +its return value becomes the body: + +```python +F.array_transform(col("a"), lambda v: v * 2) # map: [1,2,3] -> [2,4,6] +F.array_filter(col("a"), lambda v: v > 2) # filter: [1,2,3] -> [3] +F.array_any_match(col("a"), lambda v: v > 3) # predicate: any element > 3 +``` + +For explicit parameter names, build the lambda by hand: + +```python +F.array_transform(col("a"), F.lambda_(["v"], F.lambda_var("v") * lit(2))) +``` + ## SQL-to-DataFrame Reference | SQL | DataFrame API | From d021e6afa8e08bee42fb9673ba811352c464bdf7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 4 Jun 2026 07:57:47 -0400 Subject: [PATCH 51/83] feat: import user-defined physical optimizer rules over FFI (#1557) * feat: user-defined OptimizerRule and AnalyzerRule from Python Expose `SessionContext.add_optimizer_rule` and `SessionContext.add_analyzer_rule` symmetric with the existing `remove_optimizer_rule`. Each accepts a Python subclass of the new `datafusion.optimizer.OptimizerRule` / `AnalyzerRule` ABCs. Implementation: * New `crates/core/src/optimizer_rules.rs` wraps user Python instances in `PyOptimizerRuleAdapter` / `PyAnalyzerRuleAdapter`, which implement the upstream `OptimizerRule` / `AnalyzerRule` traits. * `OptimizerRule.rewrite(plan)` returns `None` for "no change" or a new `LogicalPlan`. The adapter maps that to `Transformed::no` / `Transformed::yes` so the upstream optimizer's fixed-point loop terminates correctly. * `AnalyzerRule.analyze(plan)` must always return a `LogicalPlan`; returning `None` surfaces a `DataFusionError::Execution` naming the offending rule. * The upstream `&dyn OptimizerConfig` / `&ConfigOptions` arguments are not surfaced to Python in this MVP; rules that need configuration should capture it at construction time (for example by holding a `SessionContext` reference) or be implemented in Rust. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: import FFI physical optimizer rules; drop Python logical rules Replace the Python-defined OptimizerRule/AnalyzerRule approach with FFI-imported physical optimizer rules. The Python logical-rule approach could observe plans but not transform them: there are no Python constructors for LogicalPlan node variants, so a rule could only return None or the input plan unchanged. The audience for custom rules also overlaps strongly with people who can write Rust. DataFusion exposes no FFI bridge for the logical OptimizerRule/AnalyzerRule traits, but it does export FFI_PhysicalOptimizerRule for the physical PhysicalOptimizerRule trait. This commit imports those instead. Changes: * Remove crates/core/src/optimizer_rules.rs, python/datafusion/optimizer.py, python/tests/test_optimizer.py, and the SessionContext.add_optimizer_rule / add_analyzer_rule methods. remove_optimizer_rule is unchanged (pre-existing). * New crates/core/src/physical_optimizer.rs reads a __datafusion_physical_optimizer_rule__ capsule and converts it via Arc::from(&FFI_PhysicalOptimizerRule). * SessionContext gains a physical_optimizer_rules constructor argument. Upstream offers no API to add physical rules to a live context, so they are appended to the builder at construction time only. * The datafusion-ffi-example crate gains MyPhysicalOptimizerRule, a counter-backed rule used by _test_physical_optimizer_rule.py to prove the rule fires over FFI during physical planning. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: type physical_optimizer_rules with an Exportable Protocol Replace the `list[Any]` hint on the SessionContext `physical_optimizer_rules` argument with a `PhysicalOptimizerRuleExportable` Protocol, matching the existing `TableProviderExportable` / `*Exportable` pattern used for other FFI-capsule objects. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: reference PhysicalOptimizerRuleExportable in SessionContext docstring Point the `physical_optimizer_rules` argument docs at the new `PhysicalOptimizerRuleExportable` Protocol instead of describing the duck type inline. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: move FFI capsule detail to PhysicalOptimizerRuleExportable The PyCapsule / FFI_PhysicalOptimizerRule mechanics describe the Protocol, not the SessionContext constructor. Move that detail onto PhysicalOptimizerRuleExportable and leave the constructor argument docs focused on behavior. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: drop redundant comment in SessionContext constructor Remove the explanatory comment about FFI bridge availability; the same information already lives on PhysicalOptimizerRuleExportable. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: drop module-level doc comment from physical_optimizer Sibling FFI-import modules (udf, udaf, catalog, table) carry no module-level docs, and the rst-style markup did not match Rust conventions. The function doc comment already states intent. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: import physical optimizer rule via from_pycapsule! macro Replace the hand-written crates/core/src/physical_optimizer.rs with a `from_pycapsule!` invocation in the util crate, matching `physical_codec_from_pycapsule` and the other FFI capsule importers. The macro already handles the hasattr/getattr/cast/validate/pointer_checked sequence and the infallible `Arc::from(&FFI)` conversion, so the dedicated module is no longer needed. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: trim PhysicalOptimizerRuleExportable docstring Drop the sentence about logical-rule FFI availability; it is background, not type-hint information, and keeps the Protocol docstring in line with the other *Exportable hints. Co-Authored-By: Claude Opus 4.7 (1M context) * Minor refactor * refactor: register physical optimizer rules via live add method Drop the `physical_optimizer_rules` constructor argument on `SessionContext` and replace it with `add_physical_optimizer_rule`, matching the existing `register_*` shape on the same class. The new method rebuilds the session state via `SessionStateBuilder::new_from_existing` so previously registered tables, UDFs, and catalogs are preserved. Co-Authored-By: Claude Opus 4.7 (1M context) * test: drop redundant FFI physical optimizer rule export test Coverage subsumed by test_ffi_physical_optimizer_rule_runs_during_planning, which exercises the same capsule export via add_physical_optimizer_rule. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/context.rs | 14 ++- crates/util/src/lib.rs | 9 ++ .../tests/_test_physical_optimizer_rule.py | 45 +++++++++ examples/datafusion-ffi-example/src/lib.rs | 3 + .../src/physical_optimizer.rs | 98 +++++++++++++++++++ python/datafusion/context.py | 34 +++++++ 6 files changed, 202 insertions(+), 1 deletion(-) create mode 100644 examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py create mode 100644 examples/datafusion-ffi-example/src/physical_optimizer.rs diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index d714861a6..da0df751b 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -59,7 +59,8 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ create_logical_extension_capsule, create_physical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx, get_tokio_runtime, - physical_codec_from_pycapsule, spawn_future, wait_for_future, + physical_codec_from_pycapsule, physical_optimizer_rule_from_pycapsule, spawn_future, + wait_for_future, }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; @@ -1195,6 +1196,17 @@ impl PySessionContext { self.ctx.remove_optimizer_rule(name) } + pub fn add_physical_optimizer_rule(&self, rule: Bound<'_, PyAny>) -> PyDataFusionResult<()> { + let rule = physical_optimizer_rule_from_pycapsule(&rule)?; + let state_ref = self.ctx.state_ref(); + let mut guard = state_ref.write(); + let new_state = SessionStateBuilder::new_from_existing(guard.clone()) + .with_physical_optimizer_rule(rule) + .build(); + *guard = new_state; + Ok(()) + } + pub fn table_provider(&self, name: &str, py: Python) -> PyResult { let provider = wait_for_future(py, self.ctx.table_provider(name)) // Outer error: runtime/async failure diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 75b8eeec4..07aa0a2d5 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -24,7 +24,9 @@ use datafusion::datasource::TableProvider; use datafusion::execution::TaskContext; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::Volatility; +use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion_ffi::execution::FFI_TaskContextProvider; +use datafusion_ffi::physical_optimizer::FFI_PhysicalOptimizerRule; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_ffi::table_provider::FFI_TableProvider; @@ -332,6 +334,13 @@ from_pycapsule!( dyn PhysicalExtensionCodec ); +from_pycapsule!( + physical_optimizer_rule_from_pycapsule, + "datafusion_physical_optimizer_rule", + FFI_PhysicalOptimizerRule, + dyn PhysicalOptimizerRule + Send + Sync +); + try_from_pycapsule!( task_context_from_pycapsule, "datafusion_task_context_provider", diff --git a/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py b/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py new file mode 100644 index 000000000..0c877d78e --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_physical_optimizer_rule.py @@ -0,0 +1,45 @@ +# 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. + +from __future__ import annotations + +import pyarrow as pa +from datafusion import SessionContext +from datafusion_ffi_example import MyPhysicalOptimizerRule + + +def test_ffi_physical_optimizer_rule_runs_during_planning(): + """A rule added via add_physical_optimizer_rule is invoked while the + physical plan is built, and the query still returns correct results.""" + rule = MyPhysicalOptimizerRule() + ctx = SessionContext() + ctx.add_physical_optimizer_rule(rule) + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3])], + names=["a"], + ) + ctx.register_record_batches("t", [[batch]]) + + before = rule.optimize_calls() + result = ctx.sql("SELECT a FROM t").collect() + after = rule.optimize_calls() + + assert after > before, ( + f"Expected user FFI physical optimizer rule to fire, " + f"before={before} after={after}" + ) + assert result[0].column(0).to_pylist() == [1, 2, 3] diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 3323ac982..eccf7b81a 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -22,6 +22,7 @@ use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogP use crate::config::MyConfig; use crate::logical_extension_codec::MyLogicalExtensionCodec; use crate::physical_extension_codec::MyPhysicalExtensionCodec; +use crate::physical_optimizer::MyPhysicalOptimizerRule; use crate::scalar_udf::IsNullUDF; use crate::table_function::MyTableFunction; use crate::table_provider::MyTableProvider; @@ -33,6 +34,7 @@ pub(crate) mod catalog_provider; pub(crate) mod config; pub(crate) mod logical_extension_codec; pub(crate) mod physical_extension_codec; +pub(crate) mod physical_optimizer; pub(crate) mod scalar_udf; pub(crate) mod table_function; pub(crate) mod table_provider; @@ -55,5 +57,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/physical_optimizer.rs b/examples/datafusion-ffi-example/src/physical_optimizer.rs new file mode 100644 index 000000000..0acd1bb4a --- /dev/null +++ b/examples/datafusion-ffi-example/src/physical_optimizer.rs @@ -0,0 +1,98 @@ +// 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. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_ffi::physical_optimizer::FFI_PhysicalOptimizerRule; +use datafusion_python_util::get_tokio_runtime; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +/// A physical optimizer rule that leaves every plan unchanged but bumps a +/// shared counter each time it runs. Tests use the counter to prove that a +/// session built with this rule actually routed physical planning through a +/// user-supplied [`PhysicalOptimizerRule`] over FFI. +#[derive(Debug)] +struct CountingPhysicalOptimizerRule { + optimize_calls: Arc, +} + +impl PhysicalOptimizerRule for CountingPhysicalOptimizerRule { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + self.optimize_calls.fetch_add(1, Ordering::SeqCst); + Ok(plan) + } + + fn name(&self) -> &str { + "counting_physical_optimizer_rule" + } + + fn schema_check(&self) -> bool { + // The plan is returned unchanged, so the schema is preserved. + true + } +} + +/// Python-visible handle that produces an [`FFI_PhysicalOptimizerRule`] and +/// exposes the shared call counter. +#[pyclass( + from_py_object, + name = "MyPhysicalOptimizerRule", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Default, Clone)] +pub(crate) struct MyPhysicalOptimizerRule { + optimize_calls: Arc, +} + +#[pymethods] +impl MyPhysicalOptimizerRule { + #[new] + fn new() -> Self { + Self::default() + } + + fn optimize_calls(&self) -> usize { + self.optimize_calls.load(Ordering::SeqCst) + } + + fn __datafusion_physical_optimizer_rule__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let rule: Arc = + Arc::new(CountingPhysicalOptimizerRule { + optimize_calls: Arc::clone(&self.optimize_calls), + }); + + let runtime = get_tokio_runtime().handle().clone(); + let ffi = FFI_PhysicalOptimizerRule::new(rule, Some(runtime)); + + let name = cr"datafusion_physical_optimizer_rule".into(); + PyCapsule::new(py, ffi, Some(name)) + } +} diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 52bd600c3..accb60f19 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -133,6 +133,16 @@ class TableProviderExportable(Protocol): def __datafusion_table_provider__(self, session: Any) -> object: ... # noqa: D105 +class PhysicalOptimizerRuleExportable(Protocol): + """Type hint for object that has __datafusion_physical_optimizer_rule__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``, + typically produced by a separate compiled extension. + """ + + def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 + + class SessionConfig: """Session configuration options.""" @@ -1566,6 +1576,30 @@ def remove_optimizer_rule(self, name: str) -> bool: """ return self.ctx.remove_optimizer_rule(name) + def add_physical_optimizer_rule( + self, rule: PhysicalOptimizerRuleExportable + ) -> None: + """Append a user-defined physical optimizer rule to the session. + + The rule is imported via its ``__datafusion_physical_optimizer_rule__`` + PyCapsule, typically produced by a separate compiled extension. The + underlying :class:`SessionState` is rebuilt from its current state + with the new rule appended, so previously registered tables, UDFs, + and catalogs are preserved. + + Args: + rule: Object exposing ``__datafusion_physical_optimizer_rule__``, + a :class:`PhysicalOptimizerRuleExportable`. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> from my_extension import MyPhysicalOptimizerRule # doctest: +SKIP + >>> rule = MyPhysicalOptimizerRule() # doctest: +SKIP + >>> ctx.add_physical_optimizer_rule(rule) # doctest: +SKIP + """ + self.ctx.add_physical_optimizer_rule(rule) + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. From 4a7761736f57376a9c769efd46a8b34e7f46b8f0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 5 Jun 2026 12:01:53 -0400 Subject: [PATCH 52/83] feat: expose SessionContext.copied_config and parse_capacity_limit (#1570) Adds two small additions to SessionContext that mirror upstream: - copied_config(): returns a copy of the active SessionConfig wrapped in the existing SessionConfig Python class. Useful when callers want to seed a new context from another context's settings, or inspect the current configuration without sharing mutable state. - parse_capacity_limit(config_name, limit): static helper that parses size strings like "100M", "1.5G", "512K", or "0" into a byte count. Useful when configuring a RuntimeEnvBuilder from human-friendly inputs. Wraps SessionContext::parse_capacity_limit; the deprecated parse_memory_limit is intentionally not exposed. Three other items from the same gap cluster (runtime_env, copied_table_options, the deprecated parse_memory_limit) are not included here. The first two would require wrapping new Rust types (RuntimeEnv, TableOptions) whose surface is much larger than the accessors themselves; the third is deprecated upstream. Those are filed as separate follow-up issues. Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/context.rs | 14 +++++++++++++ python/datafusion/context.py | 39 ++++++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index da0df751b..ce6629385 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1162,6 +1162,20 @@ impl PySessionContext { self.ctx.session_id() } + /// Return a copy of the active `SessionConfig`. Mutating the returned + /// config does not affect this context. + pub fn copied_config(&self) -> PySessionConfig { + self.ctx.copied_config().into() + } + + /// Parse a string like `"100M"`, `"1.5G"`, or `"512K"` into a byte count. + /// `"0"` is accepted and returns 0. Use this when constructing a + /// `RuntimeEnvBuilder` from a human-friendly size string. + #[staticmethod] + pub fn parse_capacity_limit(config_name: &str, limit: &str) -> PyDataFusionResult { + Ok(SessionContext::parse_capacity_limit(config_name, limit)?) + } + pub fn session_start_time(&self) -> String { self.ctx.session_start_time().to_rfc3339() } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index accb60f19..3be320666 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1506,6 +1506,45 @@ def enable_ident_normalization(self) -> bool: """ return self.ctx.enable_ident_normalization() + def copied_config(self) -> SessionConfig: + """Return a copy of the active :py:class:`SessionConfig`. + + Mutating the returned config does not affect this context; use + the result when you need a starting point for a new context or + want to inspect the current settings independent of further + changes here. + + Examples: + >>> ctx = SessionContext(SessionConfig().with_batch_size(1024)) + >>> isinstance(ctx.copied_config(), SessionConfig) + True + """ + config = SessionConfig() + config.config_internal = self.ctx.copied_config() + return config + + @staticmethod + def parse_capacity_limit(config_name: str, limit: str) -> int: + """Parse a size string into a byte count. + + Accepts strings like ``"100M"``, ``"1.5G"``, or ``"512K"``. + ``"0"`` is accepted and returns 0. ``config_name`` is used purely + for error messages and identifies which configuration setting the + limit belongs to. Use this helper when constructing a + :py:class:`RuntimeEnvBuilder` from a human-friendly size string. + + Examples: + >>> SessionContext.parse_capacity_limit( + ... "datafusion.runtime.memory_limit", "1M" + ... ) + 1048576 + >>> SessionContext.parse_capacity_limit( + ... "datafusion.runtime.memory_limit", "0" + ... ) + 0 + """ + return SessionContextInternal.parse_capacity_limit(config_name, limit) + def parse_sql_expr(self, sql: str, schema: DFSchema) -> Expr: """Parse a SQL expression string into a logical expression. From 23062f78ad55a7121517c5ab00742503e95c7342 Mon Sep 17 00:00:00 2001 From: Nuno Faria Date: Fri, 5 Jun 2026 17:02:21 +0100 Subject: [PATCH 53/83] fix: Skip `fork` and `forkserver` on `win32` (#1566) * fix: Skip fork and forkserver on win32 * Fix fmt --- python/tests/test_pickle_multiprocessing.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/python/tests/test_pickle_multiprocessing.py b/python/tests/test_pickle_multiprocessing.py index fcce49d97..c22b52d29 100644 --- a/python/tests/test_pickle_multiprocessing.py +++ b/python/tests/test_pickle_multiprocessing.py @@ -84,11 +84,18 @@ def _skip_if_multiprocessing_unavailable(): pytest.param( "fork", marks=pytest.mark.skipif( - sys.platform == "darwin", - reason="fork start method is unsafe with PyArrow/tokio on macOS", + sys.platform in ("darwin", "win32"), + reason="fork start method is not supported on Windows " + "and unsafe with PyArrow/tokio on macOS", + ), + ), + pytest.param( + "forkserver", + marks=pytest.mark.skipif( + sys.platform == "win32", + reason="forkserver start method is not supported on Windows", ), ), - "forkserver", "spawn", ] From bfa14f4ffa879c83acfab2f1d480d9ed474baf7d Mon Sep 17 00:00:00 2001 From: Daniel Mesejo Date: Fri, 5 Jun 2026 18:09:17 +0200 Subject: [PATCH 54/83] refactor(context): deduplicate register/read option-building logic (#1479) * refactor(context): deduplicate register/read option-building logic Extract shared helpers (convert_partition_cols, convert_file_sort_order, build_parquet/json/avro_options, convert_csv_options), standardize path types to &str, and remove redundant intermediate variables. * refactor(context): accept PathBuf for path arguments in register/read methods Change path parameters from &str to PathBuf in all register/read methods (register_listing_table, register_parquet, register_json, register_avro, register_arrow, read_json, read_parquet, read_avro, read_arrow) so callers can pass either a Python str or a pathlib.Path object. For register_csv and read_csv, which take &Bound to handle lists, extract path elements as PathBuf rather than String for the same reason. Add a path_to_str helper that converts PathBuf to &str, returning an explicit error for non-UTF-8 paths rather than silently corrupting them. Add build_arrow_options helper to deduplicate register_arrow/read_arrow option-building logic, consistent with the existing parquet/json/avro helpers. Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 --- crates/core/src/context.rs | 363 ++++++++++++++++++----------------- python/datafusion/context.py | 28 ++- 2 files changed, 196 insertions(+), 195 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index ce6629385..0db49625f 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -16,7 +16,7 @@ // under the License. use std::collections::{HashMap, HashSet}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::ptr::NonNull; use std::str::FromStr; use std::sync::Arc; @@ -477,7 +477,7 @@ impl PySessionContext { pub fn register_listing_table( &self, name: &str, - path: &str, + path: PathBuf, table_partition_cols: Vec<(String, PyArrowType)>, file_extension: &str, schema: Option>, @@ -486,20 +486,9 @@ impl PySessionContext { ) -> PyDataFusionResult<()> { let options = ListingOptions::new(Arc::new(ParquetFormat::new())) .with_file_extension(file_extension) - .with_table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ) - .with_file_sort_order( - file_sort_order - .unwrap_or_default() - .into_iter() - .map(|e| e.into_iter().map(|f| f.into()).collect()) - .collect(), - ); - let table_path = ListingTableUrl::parse(path)?; + .with_table_partition_cols(convert_partition_cols(table_partition_cols)) + .with_file_sort_order(convert_file_sort_order(file_sort_order)); + let table_path = ListingTableUrl::parse(path_to_str(&path)?)?; let resolved_schema: SchemaRef = match schema { Some(s) => Arc::new(s.0), None => { @@ -866,7 +855,7 @@ impl PySessionContext { pub fn register_parquet( &self, name: &str, - path: &str, + path: PathBuf, table_partition_cols: Vec<(String, PyArrowType)>, parquet_pruning: bool, file_extension: &str, @@ -875,25 +864,19 @@ impl PySessionContext { file_sort_order: Option>>, py: Python, ) -> PyDataFusionResult<()> { - let mut options = ParquetReadOptions::default() - .table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ) - .parquet_pruning(parquet_pruning) - .skip_metadata(skip_metadata); - options.file_extension = file_extension; - options.schema = schema.as_ref().map(|x| &x.0); - options.file_sort_order = file_sort_order - .unwrap_or_default() - .into_iter() - .map(|e| e.into_iter().map(|f| f.into()).collect()) - .collect(); - - let result = self.ctx.register_parquet(name, path, options); - wait_for_future(py, result)??; + let options = build_parquet_options( + table_partition_cols, + parquet_pruning, + file_extension, + skip_metadata, + &schema, + file_sort_order, + ); + wait_for_future( + py, + self.ctx + .register_parquet(name, path_to_str(&path)?, options), + )??; Ok(()) } @@ -907,19 +890,24 @@ impl PySessionContext { options: Option<&PyCsvReadOptions>, py: Python, ) -> PyDataFusionResult<()> { - let options = options - .map(|opts| opts.try_into()) - .transpose()? - .unwrap_or_default(); + let options = convert_csv_options(options)?; if path.is_instance_of::() { - let paths = path.extract::>()?; - let result = self.register_csv_from_multiple_paths(name, paths, options); - wait_for_future(py, result)??; + let paths = path + .extract::>()? + .iter() + .map(|p| path_to_str(p).map(str::to_owned)) + .collect::>>()?; + wait_for_future( + py, + self.register_csv_from_multiple_paths(name, paths, options), + )??; } else { - let path = path.extract::()?; - let result = self.ctx.register_csv(name, &path, options); - wait_for_future(py, result)??; + let path = path.extract::()?; + wait_for_future( + py, + self.ctx.register_csv(name, path_to_str(&path)?, options), + )??; } Ok(()) @@ -944,25 +932,17 @@ impl PySessionContext { file_compression_type: Option, py: Python, ) -> PyDataFusionResult<()> { - let path = path - .to_str() - .ok_or_else(|| PyValueError::new_err("Unable to convert path to a string"))?; - - let mut options = JsonReadOptions::default() - .file_compression_type(parse_file_compression_type(file_compression_type)?) - .table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ); - options.schema_infer_max_records = schema_infer_max_records; - options.file_extension = file_extension; - options.schema = schema.as_ref().map(|x| &x.0); - - let result = self.ctx.register_json(name, path, options); - wait_for_future(py, result)??; - + let options = build_json_options( + table_partition_cols, + file_compression_type, + schema_infer_max_records, + file_extension, + &schema, + )?; + wait_for_future( + py, + self.ctx.register_json(name, path_to_str(&path)?, options), + )??; Ok(()) } @@ -981,22 +961,11 @@ impl PySessionContext { table_partition_cols: Vec<(String, PyArrowType)>, py: Python, ) -> PyDataFusionResult<()> { - let path = path - .to_str() - .ok_or_else(|| PyValueError::new_err("Unable to convert path to a string"))?; - - let mut options = AvroReadOptions::default().table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ); - options.file_extension = file_extension; - options.schema = schema.as_ref().map(|x| &x.0); - - let result = self.ctx.register_avro(name, path, options); - wait_for_future(py, result)??; - + let options = build_avro_options(table_partition_cols, file_extension, &schema); + wait_for_future( + py, + self.ctx.register_avro(name, path_to_str(&path)?, options), + )??; Ok(()) } @@ -1004,23 +973,17 @@ impl PySessionContext { pub fn register_arrow( &self, name: &str, - path: &str, + path: PathBuf, schema: Option>, file_extension: &str, table_partition_cols: Vec<(String, PyArrowType)>, py: Python, ) -> PyDataFusionResult<()> { - let mut options = ArrowReadOptions::default().table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ); - options.file_extension = file_extension; - options.schema = schema.as_ref().map(|x| &x.0); - - let result = self.ctx.register_arrow(name, path, options); - wait_for_future(py, result)??; + let options = build_arrow_options(table_partition_cols, file_extension, &schema); + wait_for_future( + py, + self.ctx.register_arrow(name, path_to_str(&path)?, options), + )??; Ok(()) } @@ -1242,27 +1205,14 @@ impl PySessionContext { file_compression_type: Option, py: Python, ) -> PyDataFusionResult { - let path = path - .to_str() - .ok_or_else(|| PyValueError::new_err("Unable to convert path to a string"))?; - let mut options = JsonReadOptions::default() - .table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ) - .file_compression_type(parse_file_compression_type(file_compression_type)?); - options.schema_infer_max_records = schema_infer_max_records; - options.file_extension = file_extension; - let df = if let Some(schema) = schema { - options.schema = Some(&schema.0); - let result = self.ctx.read_json(path, options); - wait_for_future(py, result)?? - } else { - let result = self.ctx.read_json(path, options); - wait_for_future(py, result)?? - }; + let options = build_json_options( + table_partition_cols, + file_compression_type, + schema_infer_max_records, + file_extension, + &schema, + )?; + let df = wait_for_future(py, self.ctx.read_json(path_to_str(&path)?, options))??; Ok(PyDataFrame::new(df)) } @@ -1275,23 +1225,18 @@ impl PySessionContext { options: Option<&PyCsvReadOptions>, py: Python, ) -> PyDataFusionResult { - let options = options - .map(|opts| opts.try_into()) - .transpose()? - .unwrap_or_default(); + let options = convert_csv_options(options)?; - if path.is_instance_of::() { - let paths = path.extract::>()?; - let paths = paths.iter().map(|p| p as &str).collect::>(); - let result = self.ctx.read_csv(paths, options); - let df = PyDataFrame::new(wait_for_future(py, result)??); - Ok(df) + let paths: Vec = if path.is_instance_of::() { + path.extract::>()? + .iter() + .map(|p| path_to_str(p).map(str::to_owned)) + .collect::>()? } else { - let path = path.extract::()?; - let result = self.ctx.read_csv(path, options); - let df = PyDataFrame::new(wait_for_future(py, result)??); - Ok(df) - } + vec![path_to_str(&path.extract::()?)?.to_owned()] + }; + let df = wait_for_future(py, self.ctx.read_csv(paths, options))??; + Ok(PyDataFrame::new(df)) } #[allow(clippy::too_many_arguments)] @@ -1305,7 +1250,7 @@ impl PySessionContext { file_sort_order=None))] pub fn read_parquet( &self, - path: &str, + path: PathBuf, table_partition_cols: Vec<(String, PyArrowType)>, parquet_pruning: bool, file_extension: &str, @@ -1314,25 +1259,18 @@ impl PySessionContext { file_sort_order: Option>>, py: Python, ) -> PyDataFusionResult { - let mut options = ParquetReadOptions::default() - .table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ) - .parquet_pruning(parquet_pruning) - .skip_metadata(skip_metadata); - options.file_extension = file_extension; - options.schema = schema.as_ref().map(|x| &x.0); - options.file_sort_order = file_sort_order - .unwrap_or_default() - .into_iter() - .map(|e| e.into_iter().map(|f| f.into()).collect()) - .collect(); - - let result = self.ctx.read_parquet(path, options); - let df = PyDataFrame::new(wait_for_future(py, result)??); + let options = build_parquet_options( + table_partition_cols, + parquet_pruning, + file_extension, + skip_metadata, + &schema, + file_sort_order, + ); + let df = PyDataFrame::new(wait_for_future( + py, + self.ctx.read_parquet(path_to_str(&path)?, options), + )??); Ok(df) } @@ -1340,50 +1278,28 @@ impl PySessionContext { #[pyo3(signature = (path, schema=None, table_partition_cols=vec![], file_extension=".avro"))] pub fn read_avro( &self, - path: &str, + path: PathBuf, schema: Option>, table_partition_cols: Vec<(String, PyArrowType)>, file_extension: &str, py: Python, ) -> PyDataFusionResult { - let mut options = AvroReadOptions::default().table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ); - options.file_extension = file_extension; - let df = if let Some(schema) = schema { - options.schema = Some(&schema.0); - let read_future = self.ctx.read_avro(path, options); - wait_for_future(py, read_future)?? - } else { - let read_future = self.ctx.read_avro(path, options); - wait_for_future(py, read_future)?? - }; + let options = build_avro_options(table_partition_cols, file_extension, &schema); + let df = wait_for_future(py, self.ctx.read_avro(path_to_str(&path)?, options))??; Ok(PyDataFrame::new(df)) } #[pyo3(signature = (path, schema=None, file_extension=".arrow", table_partition_cols=vec![]))] pub fn read_arrow( &self, - path: &str, + path: PathBuf, schema: Option>, file_extension: &str, table_partition_cols: Vec<(String, PyArrowType)>, py: Python, ) -> PyDataFusionResult { - let mut options = ArrowReadOptions::default().table_partition_cols( - table_partition_cols - .into_iter() - .map(|(name, ty)| (name, ty.0)) - .collect::>(), - ); - options.file_extension = file_extension; - options.schema = schema.as_ref().map(|x| &x.0); - - let result = self.ctx.read_arrow(path, options); - let df = wait_for_future(py, result)??; + let options = build_arrow_options(table_partition_cols, file_extension, &schema); + let df = wait_for_future(py, self.ctx.read_arrow(path_to_str(&path)?, options))??; Ok(PyDataFrame::new(df)) } @@ -1523,7 +1439,7 @@ impl PySessionContext { // check if the file extension matches the expected extension for path in &table_paths { let file_path = path.as_str(); - if !file_path.ends_with(option_extension.clone().as_str()) && !path.is_collection() { + if !file_path.ends_with(option_extension.as_str()) && !path.is_collection() { return exec_err!( "File path '{file_path}' does not match the expected extension '{option_extension}'" ); @@ -1594,6 +1510,97 @@ pub fn parse_file_compression_type( }) } +fn path_to_str(path: &Path) -> PyDataFusionResult<&str> { + path.to_str() + .ok_or_else(|| PyValueError::new_err("Unable to convert path to a string").into()) +} + +fn convert_csv_options( + options: Option<&PyCsvReadOptions>, +) -> PyDataFusionResult> { + Ok(options + .map(|opts| opts.try_into()) + .transpose()? + .unwrap_or_default()) +} + +fn convert_partition_cols( + table_partition_cols: Vec<(String, PyArrowType)>, +) -> Vec<(String, DataType)> { + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect() +} + +fn convert_file_sort_order( + file_sort_order: Option>>, +) -> Vec> { + file_sort_order + .unwrap_or_default() + .into_iter() + .map(|e| e.into_iter().map(|f| f.into()).collect()) + .collect() +} + +fn build_parquet_options<'a>( + table_partition_cols: Vec<(String, PyArrowType)>, + parquet_pruning: bool, + file_extension: &'a str, + skip_metadata: bool, + schema: &'a Option>, + file_sort_order: Option>>, +) -> ParquetReadOptions<'a> { + let mut options = ParquetReadOptions::default() + .table_partition_cols(convert_partition_cols(table_partition_cols)) + .parquet_pruning(parquet_pruning) + .skip_metadata(skip_metadata); + options.file_extension = file_extension; + options.schema = schema.as_ref().map(|x| &x.0); + options.file_sort_order = convert_file_sort_order(file_sort_order); + options +} + +fn build_json_options<'a>( + table_partition_cols: Vec<(String, PyArrowType)>, + file_compression_type: Option, + schema_infer_max_records: usize, + file_extension: &'a str, + schema: &'a Option>, +) -> Result, PyErr> { + let mut options = JsonReadOptions::default() + .table_partition_cols(convert_partition_cols(table_partition_cols)) + .file_compression_type(parse_file_compression_type(file_compression_type)?); + options.schema_infer_max_records = schema_infer_max_records; + options.file_extension = file_extension; + options.schema = schema.as_ref().map(|x| &x.0); + Ok(options) +} + +fn build_arrow_options<'a>( + table_partition_cols: Vec<(String, PyArrowType)>, + file_extension: &'a str, + schema: &'a Option>, +) -> ArrowReadOptions<'a> { + let mut options = ArrowReadOptions::default() + .table_partition_cols(convert_partition_cols(table_partition_cols)); + options.file_extension = file_extension; + options.schema = schema.as_ref().map(|x| &x.0); + options +} + +fn build_avro_options<'a>( + table_partition_cols: Vec<(String, PyArrowType)>, + file_extension: &'a str, + schema: &'a Option>, +) -> AvroReadOptions<'a> { + let mut options = AvroReadOptions::default() + .table_partition_cols(convert_partition_cols(table_partition_cols)); + options.file_extension = file_extension; + options.schema = schema.as_ref().map(|x| &x.0); + options +} + impl From for SessionContext { fn from(ctx: PySessionContext) -> SessionContext { ctx.ctx.as_ref().clone() diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 3be320666..5dfeed719 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -641,7 +641,7 @@ def register_listing_table( table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_listing_table( name, - str(path), + path, table_partition_cols, file_extension, schema, @@ -1055,7 +1055,7 @@ def register_parquet( table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_parquet( name, - str(path), + path, table_partition_cols, parquet_pruning, file_extension, @@ -1097,8 +1097,6 @@ def register_csv( options: Set advanced options for CSV reading. This cannot be combined with any of the other options in this method. """ - path_arg = [str(p) for p in path] if isinstance(path, list) else str(path) - if options is not None and ( schema is not None or not has_header @@ -1132,7 +1130,7 @@ def register_csv( self.ctx.register_csv( name, - path_arg, + path, options.to_inner(), ) @@ -1167,7 +1165,7 @@ def register_json( table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_json( name, - str(path), + path, schema, schema_infer_max_records, file_extension, @@ -1198,9 +1196,7 @@ def register_avro( if table_partition_cols is None: table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) - self.ctx.register_avro( - name, str(path), schema, file_extension, table_partition_cols - ) + self.ctx.register_avro(name, path, schema, file_extension, table_partition_cols) def register_arrow( self, @@ -1279,7 +1275,7 @@ def register_arrow( table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_arrow( - name, str(path), schema, file_extension, table_partition_cols + name, path, schema, file_extension, table_partition_cols ) def register_dataset(self, name: str, dataset: pa.dataset.Dataset) -> None: @@ -1693,7 +1689,7 @@ def read_json( table_partition_cols = _convert_table_partition_cols(table_partition_cols) return DataFrame( self.ctx.read_json( - str(path), + path, schema, schema_infer_max_records, file_extension, @@ -1736,8 +1732,6 @@ def read_csv( Returns: DataFrame representation of the read CSV files """ - path_arg = [str(p) for p in path] if isinstance(path, list) else str(path) - if options is not None and ( schema is not None or not has_header @@ -1773,7 +1767,7 @@ def read_csv( return DataFrame( self.ctx.read_csv( - path_arg, + path, options.to_inner(), ) ) @@ -1816,7 +1810,7 @@ def read_parquet( file_sort_order = self._convert_file_sort_order(file_sort_order) return DataFrame( self.ctx.read_parquet( - str(path), + path, table_partition_cols, parquet_pruning, file_extension, @@ -1848,7 +1842,7 @@ def read_avro( file_partition_cols = [] file_partition_cols = _convert_table_partition_cols(file_partition_cols) return DataFrame( - self.ctx.read_avro(str(path), schema, file_partition_cols, file_extension) + self.ctx.read_avro(path, schema, file_partition_cols, file_extension) ) def read_arrow( @@ -1920,7 +1914,7 @@ def read_arrow( file_partition_cols = [] file_partition_cols = _convert_table_partition_cols(file_partition_cols) return DataFrame( - self.ctx.read_arrow(str(path), schema, file_extension, file_partition_cols) + self.ctx.read_arrow(path, schema, file_extension, file_partition_cols) ) def read_empty(self) -> DataFrame: From 407298f37d9b5ccb4439cbe77ee148bf95e52de9 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sun, 7 Jun 2026 09:19:37 -0400 Subject: [PATCH 55/83] Improve documentation site layout (#1578) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: refresh theme — pydata-sphinx-theme 0.16, top navbar, dark mode Bump pydata-sphinx-theme 0.8.0 -> 0.16 to enable the modern navbar slot API and dark/light theme switcher. Configure top navbar with logo, nav links, GitHub icon, and theme switcher in conf.py. Drop the custom docs-sidebar.html override and the layout.html block that silenced the navbar — both predate the slot API and conflict with the new theme. Strip CSS overrides that fought the old theme (--pst-header-height: 0, navbar-brand sizing) and add a dark-mode variant for the inline code color and table-stripe shading. Fix the stale github_repo ("arrow-datafusion-python" -> "datafusion-python") so future Edit-on- GitHub links resolve. Bump copyright year and project name. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: collapse navbar to section landing pages Previous structure dumped every top-level toctree entry from index.rst into the navbar, producing eight items including external URLs ("Github and Issue Tracker", "Rust's API Docs", ...) that wrapped to two lines each. Introduce user-guide/index.rst and contributor-guide/index.rst as section landing pages with nested toctrees, then point index.rst at just those two plus autoapi/index. The navbar now reads "User Guide", "Contributor Guide", "API Reference" — three single-line entries. Move the external links into the index.rst body where they're discoverable without crowding navigation. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: restore external links lost in navbar restructure Add Examples and Rust API as text links in the top navbar via the pydata-sphinx-theme external_links option. Nest the code-of-conduct link inside the Contributor Guide toctree so it appears alongside the other contributor pages. Drop the duplicate "Further reading" bullet list from the landing page now that every link has a permanent home. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: render Rust API link as docs.rs icon next to GitHub Move the Rust API docs entry from external_links to icon_links and use the fa-brands fa-rust gear mark. Now sits next to the GitHub icon in navbar_end with matching visual weight instead of a wider text link. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: render sidebar nav on landing page The default pydata-sphinx-theme sidebar-nav-bs starts at the current top-level section, so the root index — which has no parent section — ends up with an empty sidebar. The theme's layout also explicitly filters sidebar-nav-bs out of the sidebar list when suppress_sidebar_ toctree() returns true (which it does for root pages), so simply overriding sidebar-nav-bs.html in templates doesn't help. Add a sidebar-globaltoc.html template that calls Sphinx's toctree() global directly to render the full document tree, and wire it through html_sidebars under a name the theme's suppress filter doesn't strip. Landing page now shows User Guide / Contributor Guide / API Reference in the sidebar with the current section expanded on inner pages. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: render expandable chevrons in sidebar nav Switch the sidebar toctree call from toctree() to generate_toctree_html with collapse=False, so nested
    s render into the DOM for every branch. The pydata-sphinx-theme JS then wraps them in
    with fa-chevron-down toggles, matching the datafusion-comet sidebar where each section with children can be expanded inline. show_nav_level=1 keeps deeper levels collapsed on first load. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: expand sidebar to show level 2 entries by default Bump show_nav_level 1 -> 2 so the landing-page sidebar opens with User Guide / Contributor Guide / API Reference already expanded to their immediate children. Deeper levels remain collapsed behind chevrons so the sidebar stays scannable. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: add Links sidebar section for external references Restore the "Links" sidebar heading that the previous site had — GitHub and Issue Tracker, Rust API Docs, Code of Conduct, Examples. Implemented as a second hidden toctree with :caption: Links so the pydata-sphinx-theme sidebar renders the heading above the four external URLs. Drop Code of Conduct from the Contributor Guide toctree since it now lives under Links instead. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: consolidate external URLs into a single Links nav item Replace the second hidden toctree (which expanded each external URL into its own navbar entry) with a dedicated links.rst landing page, and add a single "links" entry to the main toctree. Top navbar now shows User Guide / Contributor Guide / API Reference / Links — four items, no wrapping. Clicking Links opens the page that lists GitHub, Rust API Docs, Code of Conduct, and Examples. Drop the external_links Examples entry from conf.py since the same URL now lives on the Links page. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: add favicon matching the main datafusion site Drop in the same favicon.svg the main datafusion.apache.org site uses (just the Apache DataFusion mark, no wordmark) and wire it through html_favicon. Browsers and bookmarks now show the project icon instead of the generic Sphinx page glyph. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: address Copilot review feedback on sidebar config Two small follow-ups from the Copilot reviewer on #1578: - Append .html to the html_sidebars entry. Sphinx's Jinja loader resolves both "sidebar-globaltoc" and "sidebar-globaltoc.html" to the same template, but the explicit form is closer to the spelling in the Sphinx docs and is harder to misread. - Update the inline comment in sidebar-globaltoc.html that still claimed show_nav_level=1 after we bumped it to 2 in conf.py. Now describes the variable wiring instead of hard-coding a number that has to be kept in sync with conf.py. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- docs/source/_static/favicon.svg | 10 ++++ docs/source/_static/theme_overrides.css | 46 +++------------ docs/source/_templates/docs-sidebar.html | 19 ------ docs/source/_templates/layout.html | 4 -- docs/source/_templates/sidebar-globaltoc.html | 30 ++++++++++ docs/source/conf.py | 41 ++++++++++--- docs/source/contributor-guide/index.rst | 28 +++++++++ docs/source/index.rst | 44 ++------------ docs/source/links.rst | 30 ++++++++++ docs/source/user-guide/index.rst | 38 ++++++++++++ pyproject.toml | 3 +- uv.lock | 59 +++++++++++++++++-- 12 files changed, 240 insertions(+), 112 deletions(-) create mode 100644 docs/source/_static/favicon.svg delete mode 100644 docs/source/_templates/docs-sidebar.html create mode 100644 docs/source/_templates/sidebar-globaltoc.html create mode 100644 docs/source/contributor-guide/index.rst create mode 100644 docs/source/links.rst create mode 100644 docs/source/user-guide/index.rst diff --git a/docs/source/_static/favicon.svg b/docs/source/_static/favicon.svg new file mode 100644 index 000000000..bf174719b --- /dev/null +++ b/docs/source/_static/favicon.svg @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/docs/source/_static/theme_overrides.css b/docs/source/_static/theme_overrides.css index aaa40fba2..661454b12 100644 --- a/docs/source/_static/theme_overrides.css +++ b/docs/source/_static/theme_overrides.css @@ -21,62 +21,34 @@ /* Customizing with theme CSS variables */ :root { - --pst-color-active-navigation: 215, 70, 51; --pst-color-link-hover: 215, 70, 51; --pst-color-headerlink: 215, 70, 51; - /* Use normal text color (like h3, ..) instead of primary color */ - --pst-color-h1: var(--color-text-base); - --pst-color-h2: var(--color-text-base); - /* Use softer blue from bootstrap's default info color */ + /* Softer blue from bootstrap's default info color */ --pst-color-info: 23, 162, 184; - --pst-header-height: 0px; } code { color: rgb(215, 70, 51); } -.footer { - text-align: center; -} - -/* Ensure the logo is properly displayed */ - -.navbar-brand { - height: auto; - width: auto; +html[data-theme="dark"] code { + color: rgb(255, 138, 117); } -a.navbar-brand img { - height: auto; - width: auto; - max-height: 15vh; - max-width: 100%; +.footer { + text-align: center; } -/* This is the bootstrap CSS style for "table-striped". Since the theme does -not yet provide an easy way to configure this globally, it easier to simply -include this snippet here than updating each table in all rst files to -add ":class: table-striped" */ +/* Bootstrap "table-striped" applied globally so individual tables in + user-guide pages don't need ":class: table-striped" added one by one. */ .table tbody tr:nth-of-type(odd) { background-color: rgba(0, 0, 0, 0.05); } - -/* Limit the max height of the sidebar navigation section. Because in our -custimized template, there is more content above the navigation, i.e. -larger logo: if we don't decrease the max-height, it will overlap with -the footer. -Details: min(15vh, 110px) for the logo size, 8rem for search box etc*/ - -@media (min-width:720px) { - @supports (position:-webkit-sticky) or (position:sticky) { - .bd-links { - max-height: calc(100vh - min(15vh, 110px) - 8rem) - } - } +html[data-theme="dark"] .table tbody tr:nth-of-type(odd) { + background-color: rgba(255, 255, 255, 0.05); } diff --git a/docs/source/_templates/docs-sidebar.html b/docs/source/_templates/docs-sidebar.html deleted file mode 100644 index 44deeed25..000000000 --- a/docs/source/_templates/docs-sidebar.html +++ /dev/null @@ -1,19 +0,0 @@ - - - - - -
    - - -
    - - diff --git a/docs/source/_templates/layout.html b/docs/source/_templates/layout.html index 9f7880049..d83d283c7 100644 --- a/docs/source/_templates/layout.html +++ b/docs/source/_templates/layout.html @@ -1,9 +1,5 @@ {% extends "pydata_sphinx_theme/layout.html" %} -{# Silence the navbar #} -{% block docs_navbar %} -{% endblock %} - diff --git a/docs/source/_templates/sidebar-globaltoc.html b/docs/source/_templates/sidebar-globaltoc.html new file mode 100644 index 000000000..f4aa2051f --- /dev/null +++ b/docs/source/_templates/sidebar-globaltoc.html @@ -0,0 +1,30 @@ +{# Renders the global document toctree on every page (including the + landing page) with pydata-sphinx-theme's collapsible chevrons. + + The stock sidebar-nav-bs.html starts at the current section and is + stripped from the sidebar list by suppress_sidebar_toctree() on the + root page (no parent section). Using generate_toctree_html with + startdepth=0 renders the whole tree from root with the bootstrap + classes the theme's JS uses for expand/collapse toggles. Naming the + template "sidebar-globaltoc" sidesteps the suppress filter, which + matches on "sidebar-nav-bs.html" specifically. #} + diff --git a/docs/source/conf.py b/docs/source/conf.py index b2e9bb8c3..bb1473546 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -35,8 +35,8 @@ # -- Project information ----------------------------------------------------- -project = "Apache Arrow DataFusion" -copyright = "2019-2024, Apache Software Foundation" +project = "Apache DataFusion in Python" +copyright = "2019-2026, Apache Software Foundation" author = "Apache Software Foundation" @@ -115,13 +115,40 @@ def setup(sphinx) -> None: # html_theme = "pydata_sphinx_theme" -html_theme_options = {"use_edit_page_button": False, "show_toc_level": 2} +html_theme_options = { + "use_edit_page_button": False, + "show_toc_level": 2, + "logo": { + "image_light": "_static/images/original.svg", + "image_dark": "_static/images/original.svg", + "alt_text": "Apache DataFusion in Python", + }, + "navbar_start": ["navbar-logo"], + "navbar_center": ["navbar-nav"], + "navbar_end": ["navbar-icon-links", "theme-switcher"], + "icon_links": [ + { + "name": "GitHub", + "url": "https://github.com/apache/datafusion-python", + "icon": "fa-brands fa-github", + }, + { + "name": "Rust API docs (docs.rs)", + "url": "https://docs.rs/datafusion/latest/datafusion/", + "icon": "fa-brands fa-rust", + }, + ], + "secondary_sidebar_items": [], + "collapse_navigation": True, + "show_nav_level": 2, +} html_context = { "github_user": "apache", - "github_repo": "arrow-datafusion-python", + "github_repo": "datafusion-python", "github_version": "main", "doc_path": "docs/source", + "default_mode": "auto", } # Add any paths that contain custom static files (such as style sheets) here, @@ -129,16 +156,16 @@ def setup(sphinx) -> None: # so a file named "default.css" will overwrite the builtin "default.css". html_static_path = ["_static"] +html_favicon = "_static/favicon.svg" + # 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"] html_sidebars = { - "**": ["docs-sidebar.html"], + "**": ["sidebar-globaltoc.html"], } # tell myst_parser to auto-generate anchor links for headers h1, h2, h3 diff --git a/docs/source/contributor-guide/index.rst b/docs/source/contributor-guide/index.rst new file mode 100644 index 000000000..b32e08878 --- /dev/null +++ b/docs/source/contributor-guide/index.rst @@ -0,0 +1,28 @@ +.. 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. + +================= +Contributor Guide +================= + +Guides for contributors to the DataFusion in Python project. + +.. toctree:: + :maxdepth: 2 + + introduction + ffi diff --git a/docs/source/index.rst b/docs/source/index.rst index 7edb69807..6b72537da 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -52,47 +52,11 @@ Example df.show() -.. _toc.links: .. toctree:: :hidden: :maxdepth: 1 - :caption: LINKS - Github and Issue Tracker - Rust's API Docs - Code of conduct - Examples - -.. _toc.guide: -.. toctree:: - :hidden: - :maxdepth: 1 - :caption: USER GUIDE - - user-guide/introduction - user-guide/basics - user-guide/data-sources - user-guide/dataframe/index - user-guide/common-operations/index - user-guide/io/index - user-guide/configuration - user-guide/distributing-work - user-guide/sql - user-guide/upgrade-guides - user-guide/ai-coding-assistants - - -.. _toc.contributor_guide: -.. toctree:: - :hidden: - :maxdepth: 1 - :caption: CONTRIBUTOR GUIDE - - contributor-guide/introduction - contributor-guide/ffi - -.. _toc.api: -.. toctree:: - :hidden: - :maxdepth: 1 - :caption: API + user-guide/index + contributor-guide/index + API Reference + links diff --git a/docs/source/links.rst b/docs/source/links.rst new file mode 100644 index 000000000..10473f31b --- /dev/null +++ b/docs/source/links.rst @@ -0,0 +1,30 @@ +.. 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. + +===== +Links +===== + +External resources for the DataFusion in Python project. + +.. toctree:: + :maxdepth: 1 + + GitHub and Issue Tracker + Rust API Docs + Code of Conduct + Examples diff --git a/docs/source/user-guide/index.rst b/docs/source/user-guide/index.rst new file mode 100644 index 000000000..2d6b94392 --- /dev/null +++ b/docs/source/user-guide/index.rst @@ -0,0 +1,38 @@ +.. 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. + +========== +User Guide +========== + +The user guide walks through installing DataFusion in Python, building queries +with the DataFrame API or SQL, reading and writing data, and tuning execution. + +.. toctree:: + :maxdepth: 2 + + introduction + basics + data-sources + dataframe/index + common-operations/index + io/index + configuration + distributing-work + sql + upgrade-guides + ai-coding-assistants diff --git a/pyproject.toml b/pyproject.toml index 2b6a976db..e18c1d57c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -219,8 +219,9 @@ docs = [ "myst-parser>=3.0.1", "pandas>=2.0.3", "pickleshare>=0.7.5", - "pydata-sphinx-theme==0.8.0", + "pydata-sphinx-theme>=0.16,<0.17", "setuptools>=75.3.0", "sphinx-autoapi>=3.4.0", + "sphinx-reredirects>=0.1.5", "sphinx>=7.1.2", ] diff --git a/uv.lock b/uv.lock index 6673b7fe2..89617aed0 100644 --- a/uv.lock +++ b/uv.lock @@ -8,6 +8,18 @@ resolution-markers = [ "python_full_version < '3.11'", ] +[[package]] +name = "accessible-pygments" +version = "0.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bc/c1/bbac6a50d02774f91572938964c582fff4270eee73ab822a4aeea4d8b11b/accessible_pygments-0.0.5.tar.gz", hash = "sha256:40918d3e6a2b619ad424cb91e556bd3bd8865443d9f22f1dcdf79e33c8046872", size = 1377899, upload-time = "2024-05-10T11:23:10.216Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/3f/95338030883d8c8b91223b4e21744b04d11b161a3ef117295d8241f50ab4/accessible_pygments-0.0.5-py3-none-any.whl", hash = "sha256:88ae3211e68a1d0b011504b2ffc1691feafce124b845bd072ab6f9f66f34d4b7", size = 1395903, upload-time = "2024-05-10T11:23:08.421Z" }, +] + [[package]] name = "alabaster" version = "1.0.0" @@ -357,6 +369,8 @@ docs = [ { name = "setuptools" }, { name = "sphinx" }, { name = "sphinx-autoapi" }, + { name = "sphinx-reredirects", version = "0.1.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "sphinx-reredirects", version = "1.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] release = [ { name = "pygithub" }, @@ -393,10 +407,11 @@ docs = [ { name = "myst-parser", specifier = ">=3.0.1" }, { name = "pandas", specifier = ">=2.0.3" }, { name = "pickleshare", specifier = ">=0.7.5" }, - { name = "pydata-sphinx-theme", specifier = "==0.8.0" }, + { name = "pydata-sphinx-theme", specifier = ">=0.16,<0.17" }, { name = "setuptools", specifier = ">=75.3.0" }, { name = "sphinx", specifier = ">=7.1.2" }, { name = "sphinx-autoapi", specifier = ">=3.4.0" }, + { name = "sphinx-reredirects", specifier = ">=0.1.5" }, ] release = [{ name = "pygithub", specifier = "==2.5.0" }] @@ -1142,16 +1157,20 @@ wheels = [ [[package]] name = "pydata-sphinx-theme" -version = "0.8.0" +version = "0.16.1" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "accessible-pygments" }, + { name = "babel" }, { name = "beautifulsoup4" }, { name = "docutils" }, + { name = "pygments" }, { name = "sphinx" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fc/d6/3921de802cf1ee771f0e76c9068b52498aeb8eeec6b830ff931c81c7ecf3/pydata_sphinx_theme-0.8.0.tar.gz", hash = "sha256:9f72015d9c572ea92e3007ab221a8325767c426783b6b9941813e65fa988dc90", size = 1123746, upload-time = "2022-01-15T19:25:25.712Z" } +sdist = { url = "https://files.pythonhosted.org/packages/00/20/bb50f9de3a6de69e6abd6b087b52fa2418a0418b19597601605f855ad044/pydata_sphinx_theme-0.16.1.tar.gz", hash = "sha256:a08b7f0b7f70387219dc659bff0893a7554d5eb39b59d3b8ef37b8401b7642d7", size = 2412693, upload-time = "2024-12-17T10:53:39.537Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/91/26/0694318d46c7d90ab602ae27b24431e939f1600f9a4c69d1e727ec57289f/pydata_sphinx_theme-0.8.0-py3-none-any.whl", hash = "sha256:fbcbb833a07d3ad8dd997dd40dc94da18d98b41c68123ab0182b58fe92271204", size = 3284997, upload-time = "2022-01-15T19:25:23.807Z" }, + { url = "https://files.pythonhosted.org/packages/e2/0d/8ba33fa83a7dcde13eb3c1c2a0c1cc29950a048bfed6d9b0d8b6bd710b4c/pydata_sphinx_theme-0.16.1-py3-none-any.whl", hash = "sha256:225331e8ac4b32682c18fcac5a57a6f717c4e632cea5dd0e247b55155faeccde", size = 6723264, upload-time = "2024-12-17T10:53:35.645Z" }, ] [[package]] @@ -1459,6 +1478,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/d6/f2acdc2567337fd5f5dc091a4e58d8a0fb14927b9779fc1e5ecee96d9824/sphinx_autoapi-3.4.0-py3-none-any.whl", hash = "sha256:4027fef2875a22c5f2a57107c71641d82f6166bf55beb407a47aaf3ef14e7b92", size = 34095, upload-time = "2024-11-30T01:09:17.272Z" }, ] +[[package]] +name = "sphinx-reredirects" +version = "0.1.6" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "sphinx", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/16/6b/bcca2785de4071f604a722444d4d7ba8a9d40de3c14ad52fce93e6d92694/sphinx_reredirects-0.1.6.tar.gz", hash = "sha256:c491cba545f67be9697508727818d8626626366245ae64456fe29f37e9bbea64", size = 7080, upload-time = "2025-03-22T10:52:30.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/6f/0b3625be30a1a50f9e4c2cb2ec147b08f15ed0e9f8444efcf274b751300b/sphinx_reredirects-0.1.6-py3-none-any.whl", hash = "sha256:efd50c766fbc5bf40cd5148e10c00f2c00d143027de5c5e48beece93cc40eeea", size = 5675, upload-time = "2025-03-22T10:52:29.113Z" }, +] + +[[package]] +name = "sphinx-reredirects" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version >= '3.12' and python_full_version < '3.14'", + "python_full_version == '3.11.*'", +] +dependencies = [ + { name = "sphinx", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1b/8d/0e39fe2740d7d71417edf9a6424aa80ca2c27c17fc21282cdc39f90d5a40/sphinx_reredirects-1.1.0.tar.gz", hash = "sha256:fb9b195335ab14b43f8273287d0c7eeb637ba6c56c66581c11b47202f6718b29", size = 614624, upload-time = "2025-12-22T08:28:02.792Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/81/b5dd07067f3daac6d23687ec737b2d593740671ebcd145830c8f92d381c5/sphinx_reredirects-1.1.0-py3-none-any.whl", hash = "sha256:4b5692273c72cd2d4d917f4c6f87d5919e4d6114a752d4be033f7f5f6310efd9", size = 6351, upload-time = "2025-12-22T08:27:59.724Z" }, +] + [[package]] name = "sphinxcontrib-applehelp" version = "2.0.0" From 43df9f7b954c00cc40bd257317926eddb6d4092e Mon Sep 17 00:00:00 2001 From: kosiew Date: Sun, 7 Jun 2026 21:22:06 +0800 Subject: [PATCH 56/83] feat(dataframe): update group_by to accept None and normalize to empty list (#1581) - Updated `group_by` method to accept `None` and normalize it to an empty list. - Improved docstring for clarity. - Added regression test in `test_dataframe.py` to verify that `None` equals an empty list. - Updated documentation to mention that `group_by=None` is now supported. --- .../common-operations/aggregations.rst | 4 +-- python/datafusion/dataframe.py | 29 +++++++++++-------- python/tests/test_dataframe.py | 6 ++++ 3 files changed, 25 insertions(+), 14 deletions(-) diff --git a/docs/source/user-guide/common-operations/aggregations.rst b/docs/source/user-guide/common-operations/aggregations.rst index 8f218abd8..b1e43a32f 100644 --- a/docs/source/user-guide/common-operations/aggregations.rst +++ b/docs/source/user-guide/common-operations/aggregations.rst @@ -41,8 +41,8 @@ to form a single summary value. For performing an aggregation, DataFusion provid f.approx_median(col_speed).alias("Median Speed"), f.approx_percentile_cont(col_speed, 0.9).alias("90% Speed")]) -When the :code:`group_by` list is empty the aggregation is done over the whole :class:`.DataFrame`. -For grouping the :code:`group_by` list must contain at least one column. +When :code:`group_by` is :code:`None` or an empty list, the aggregation is done over the whole +:class:`.DataFrame`. For grouping the :code:`group_by` list must contain at least one column. .. ipython:: python diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index 9ac8293d6..de00ff474 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -798,7 +798,7 @@ def with_column_renamed(self, old_name: str, new_name: str) -> DataFrame: def aggregate( self, - group_by: Sequence[Expr | str] | Expr | str, + group_by: Sequence[Expr | str] | Expr | str | None, aggs: Sequence[Expr] | Expr, ) -> DataFrame: """Aggregates the rows of the current DataFrame. @@ -816,23 +816,24 @@ def aggregate( Args: group_by: Sequence of expressions or column names to group - by. A :py:class:`~datafusion.expr.GroupingSet` - expression may be included to produce multiple grouping - levels (rollup, cube, or explicit grouping sets). + by, or ``None`` for aggregation over the whole DataFrame. + A :py:class:`~datafusion.expr.GroupingSet` expression may + be included to produce multiple grouping levels (rollup, + cube, or explicit grouping sets). aggs: Sequence of expressions to aggregate. Returns: DataFrame after aggregation. Examples: - Aggregate without grouping — an empty ``group_by`` produces a - single row: + Aggregate without grouping — ``None`` or an empty ``group_by`` + produces a single row: >>> ctx = dfn.SessionContext() >>> df = ctx.from_pydict( ... {"team": ["x", "x", "y"], "score": [1, 2, 5]} ... ) - >>> df.aggregate([], [F.sum(col("score")).alias("total")]).to_pydict() + >>> df.aggregate(None, [F.sum(col("score")).alias("total")]).to_pydict() {'total': [8]} Group by a column and produce one row per group: @@ -842,11 +843,15 @@ def aggregate( ... ).sort("team").to_pydict() {'team': ['x', 'y'], 'total': [3, 5]} """ - group_by_list = ( - list(group_by) - if isinstance(group_by, Sequence) and not isinstance(group_by, Expr | str) - else [group_by] - ) + if group_by is None: + group_by_list = [] + else: + group_by_list = ( + list(group_by) + if isinstance(group_by, Sequence) + and not isinstance(group_by, Expr | str) + else [group_by] + ) aggs_list = ( list(aggs) if isinstance(aggs, Sequence) and not isinstance(aggs, Expr) diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index ab3992a79..bb21a3974 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -475,6 +475,12 @@ def test_aggregate_tuple_group_by(df): assert result_tuple == result_list +def test_aggregate_none_group_by_equivalent_to_empty_list(df): + result_none = df.aggregate(None, [f.count()]).to_pydict() + result_empty = df.aggregate([], [f.count()]).to_pydict() + assert result_none == result_empty + + def test_aggregate_tuple_aggs(df): result_list = df.aggregate("a", [f.count()]).sort("a").to_pydict() result_tuple = df.aggregate("a", (f.count(),)).sort("a").to_pydict() From 43395a46aa2597b8090248ecd9a10c601ebec0bf Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 12 Jun 2026 16:17:20 +0200 Subject: [PATCH 57/83] Remove patch now that 54 is released (#1588) --- Cargo.lock | 108 +++++++++++++++++++++++++++++++++++------------------ Cargo.toml | 9 ----- 2 files changed, 72 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d3cedb628..9d7ed77bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -791,7 +791,8 @@ dependencies = [ [[package]] name = "datafusion" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "997a31e15872606a49478e670c58302094c97cb96abb0a7d60720f8e92170040" dependencies = [ "arrow", "arrow-schema", @@ -844,7 +845,8 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7dd61161508f8f5fa1107774ea687bd753c22d83a32eebf963549f89de14139" dependencies = [ "arrow", "async-trait", @@ -868,7 +870,8 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "897c70f871277f9ce99aa38347be0d679bbe3e617156c4d2a8378cec8a2a0891" dependencies = [ "arrow", "async-trait", @@ -890,7 +893,8 @@ dependencies = [ [[package]] name = "datafusion-common" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c9ded5d87d9172319e006f2afdb9928d72dbacd6a90a458d8acb1e3b43a65" dependencies = [ "arrow", "arrow-ipc", @@ -915,7 +919,8 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "981b9dae74f78ee3d9f714fb49b01919eab975461b56149510c3ba9ea11287d1" dependencies = [ "futures", "log", @@ -925,7 +930,8 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffd7d295b2ec7c00d8a56562f41ed41062cf0af75549ed891c12a0a09eddfefe" dependencies = [ "arrow", "async-compression", @@ -960,7 +966,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "552b0b3f342f7ec41b3fbd70f6339dc82a30cfd0349e7f280e7852528085349f" dependencies = [ "arrow", "arrow-ipc", @@ -983,7 +990,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb517d08967d536284ce70afb5fe8583133779249f2d7b90587d339741a7f195" dependencies = [ "arrow", "arrow-avro", @@ -1001,7 +1009,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68850aa426b897e879c8b87e512ea8124f1d0a2869a4e51808ddaaddf1bc0ada" dependencies = [ "arrow", "async-trait", @@ -1023,7 +1032,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "402f93242ae08ef99139ee2c528a49d087efe88d5c7b2c3ff5480855a40ce54f" dependencies = [ "arrow", "async-trait", @@ -1045,7 +1055,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffd2499c1bee0eeccf6a57156105700eeeb17bc701899ac719183c4e74231450" dependencies = [ "arrow", "async-trait", @@ -1075,12 +1086,14 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb9e7e5d11130c48c8bd4e80c79a9772dd28ce6dc330baca9246205d245b9e2e" [[package]] name = "datafusion-execution" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37a8643ab852eb68864e1b72ae789e8066282dce48eea6347ffb0aee33d1ccc0" dependencies = [ "arrow", "arrow-buffer", @@ -1101,7 +1114,8 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6932f4d71eed9c8d9341476a2b845aadfabde5495d08dbcd8fc23881f49fa7a0" dependencies = [ "arrow", "arrow-schema", @@ -1123,7 +1137,8 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0225491839a31b1f7d2cb8092c2d50792e2fe1c1724e4e6d08e011f5feaf4ed2" dependencies = [ "arrow", "datafusion-common", @@ -1134,7 +1149,8 @@ dependencies = [ [[package]] name = "datafusion-ffi" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5660e8fa79fd51e29ce46f3026b67317ef738ebd633e106beb1a1907a406152" dependencies = [ "arrow", "arrow-schema", @@ -1188,7 +1204,8 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14872c47bfc3d21e53ec82f57074e6987a15941c1e2f43cde4ac6ae2746634e3" dependencies = [ "arrow", "arrow-buffer", @@ -1219,7 +1236,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75a2ca14e1b609be21e657e2d3130b2f446456b08393b377bb721a33952d2e09" dependencies = [ "arrow", "datafusion-common", @@ -1239,7 +1257,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ece74ba09092d2ef9c9b54a38445450aea292a1f8b04faf531936b723a24b3c" dependencies = [ "arrow", "datafusion-common", @@ -1250,7 +1269,8 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f3e3f9ee8ca59bf70518802107de6f1b88a9509efdc629fadc5de9d6b2d5ef5" dependencies = [ "arrow", "arrow-ord", @@ -1274,7 +1294,8 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89161dffc22cf2b50f9f4b1bee83b5221d3b4ed7c2e37fd7aa2b22a5297b3a26" dependencies = [ "arrow", "async-trait", @@ -1289,7 +1310,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7339345b226b3874037708bf5023ba1c2de705128f8457a095aae5ae9cb9c78" dependencies = [ "arrow", "datafusion-common", @@ -1305,7 +1327,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa84836dc2392df6f43d6a29d37fb56a8ebdc8b3f4e10ae8dc15861fd20278fb" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1314,7 +1337,8 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "587164e03ad68732aa9e7bfe5686e3f25970d4c64fd4bd80790749840892dae5" dependencies = [ "datafusion-doc", "quote", @@ -1324,7 +1348,8 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77f20e8cf9e8654d92f4c16b24c487353ee5bf153ffc12d5772cd399ab8cd281" dependencies = [ "arrow", "chrono", @@ -1343,7 +1368,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f015a4a82f6f7ff7e1d8d4bf3870a936752fa38b17705dfcc14adef95aa8922c" dependencies = [ "arrow", "datafusion-common", @@ -1364,7 +1390,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51e6ffff8acdfe54e0ea15ccf38115c4a9184433b0439f42907637928d00a235" dependencies = [ "arrow", "datafusion-common", @@ -1378,7 +1405,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7967a3e171c6a4bf09474b3f7a14f1a3db13ed1714ba12156f33fcce2bba54e8" dependencies = [ "arrow", "chrono", @@ -1394,7 +1422,8 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ff803e2a96054cb6d83f35f9e60fd4f42eac515e1932bd1b2dbc91d5fcbf36" dependencies = [ "arrow", "datafusion-common", @@ -1412,7 +1441,8 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "776ee54d47d15bdb126452f9ca17b03761e3b004682914beaedd3f86eb507fbc" dependencies = [ "arrow", "arrow-data", @@ -1444,7 +1474,8 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd15a1ba5d3af93808241065c6c44dbca8296a189845e8a587c45c07bf0ffae" dependencies = [ "arrow", "chrono", @@ -1470,7 +1501,8 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90042982cf9462eb06a0b81f92efa4188dae871e7ea3ab8dc61aa9c9349b2530" dependencies = [ "arrow", "datafusion-common", @@ -1480,7 +1512,8 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5fb9e5774660aa69c3ba93c610f175f75b65cb8c3776edb3626de8f3a4f4ee3" dependencies = [ "arrow", "datafusion-common", @@ -1539,7 +1572,8 @@ dependencies = [ [[package]] name = "datafusion-session" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15ce715fa2a61f4623cc234bcc14a3ef6a91f189128d5b14b468a6a17cdfc417" dependencies = [ "async-trait", "datafusion-common", @@ -1552,7 +1586,8 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6094ad36a3ed6d7ac87b20b479b2d0b118250f66cf997603829fdc65b44a7099" dependencies = [ "arrow", "bigdecimal", @@ -1570,7 +1605,8 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "54.0.0" -source = "git+https://github.com/apache/datafusion?rev=1321d60cc37ee487d1e7ce7f501357c3236b2542#1321d60cc37ee487d1e7ce7f501357c3236b2542" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b22c8f8c72d317e54fad6f85c0ef6d1e1da53cc7faadc7eea8daf0f8d86d4f2" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index e72c22368..6a3ad3ab4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,12 +71,3 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] -datafusion = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-common = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } -datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "1321d60cc37ee487d1e7ce7f501357c3236b2542" } From cd7506ad645bf7f3c01796cbb16ec5821d56d389 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 12 Jun 2026 16:45:40 +0200 Subject: [PATCH 58/83] build(deps): batch dependabot dependency updates (#1589) Combine open Dependabot updates into a single PR. GitHub Actions: - github/codeql-action 4.35.4 -> 4.36.2 (#1585) - astral-sh/setup-uv 8.1.0 -> 8.2.0 (#1584) Cargo: - chrono 0.4.44 -> 0.4.45 (#1583) - log 0.4.30 -> 0.4.32 (#1582) - uuid 1.23.1 -> 1.23.2 (#1565) Python (uv.lock): - pyarrow 22.0.0 -> 23.0.1 (#1580) - idna 3.10 -> 3.15 (#1552) - pytest 8.3.4 -> 9.0.3 (#1542) - pyjwt 2.10.1 -> 2.12.0 (#1540) - pygments 2.19.1 -> 2.20.0 (#1539) - requests 2.32.3 -> 2.33.0 (#1538) - urllib3 2.3.0 -> 2.7.0 (#1537) - pynacl 1.5.0 -> 1.6.2 (#1536) - cryptography 44.0.0 -> 46.0.7 (#1535) Co-authored-by: Claude Opus 4.8 (1M context) --- .github/workflows/build.yml | 14 +- .github/workflows/codeql.yml | 4 +- .github/workflows/test.yml | 2 +- Cargo.lock | 14 +- uv.lock | 446 +++++++++++++++++++++-------------- 5 files changed, 289 insertions(+), 191 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 593a343e1..433fc3aca 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -69,7 +69,7 @@ jobs: with: python-version: "3.12" - - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true @@ -113,7 +113,7 @@ jobs: steps: - uses: actions/checkout@v6 - - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true @@ -159,7 +159,7 @@ jobs: with: key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} - - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true @@ -238,7 +238,7 @@ jobs: with: key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} - - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true @@ -307,7 +307,7 @@ jobs: python-version: ${{ matrix.python-tag }} freethreaded: true - - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true @@ -389,7 +389,7 @@ jobs: python-version: ${{ matrix.python-tag }} freethreaded: true - - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true @@ -518,7 +518,7 @@ jobs: python-version: "3.10" - name: Install dependencies - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index e71ea6bed..2d0f166ba 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -44,11 +44,11 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@68bde559dea0fdcac2102bfdf6230c5f70eb485e # v4 + uses: github/codeql-action/init@8aad20d150bbac5944a9f9d289da16a4b0d87c1e # v4 with: languages: actions - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@68bde559dea0fdcac2102bfdf6230c5f70eb485e # v4 + uses: github/codeql-action/analyze@8aad20d150bbac5944a9f9d289da16a4b0d87c1e # v4 with: category: "/language:actions" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0c8fa4f79..bd24ee508 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -63,7 +63,7 @@ jobs: key: cargo-cache-stable-${{ hashFiles('Cargo.lock') }} - name: Install dependencies - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: enable-cache: true diff --git a/Cargo.lock b/Cargo.lock index 9d7ed77bc..4a063f749 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -571,9 +571,9 @@ dependencies = [ [[package]] name = "chrono" -version = "0.4.44" +version = "0.4.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" dependencies = [ "iana-time-zone", "num-traits", @@ -2425,9 +2425,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.30" +version = "0.4.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "616ec5685824bcc94416c6d4a7a446eea774a31efd7062c8480ba6fd06d7a6e5" +checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" [[package]] name = "lru-slab" @@ -3796,7 +3796,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.2", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys 0.61.2", @@ -4155,9 +4155,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.23.1" +version = "1.23.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" +checksum = "d258b83ceec21034727ecee8c382cfa6c3e133699b0742c64571814fb420c9f7" dependencies = [ "getrandom 0.4.2", "js-sys", diff --git a/uv.lock b/uv.lock index 89617aed0..b90f9ff11 100644 --- a/uv.lock +++ b/uv.lock @@ -34,7 +34,8 @@ name = "arro3-core" version = "0.6.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.12'" }, + { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/2a/01/f06342d2eb822153f63d188153e41fbeabb29b48247f7a11ce76c538f7d1/arro3_core-0.6.5.tar.gz", hash = "sha256:768078887cd7ac82de4736f94bbd91f6d660f10779848bd5b019f511badd9d75", size = 107522, upload-time = "2025-10-13T23:12:38.872Z" } wheels = [ @@ -96,7 +97,7 @@ name = "astroid" version = "3.3.8" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/80/c5/5c83c48bbf547f3dd8b587529db7cf5a265a3368b33e85e76af8ff6061d3/astroid-3.3.8.tar.gz", hash = "sha256:a88c7994f914a4ea8572fac479459f4955eeccc877be3f2d959a33273b0cf40b", size = 398196, upload-time = "2024-12-24T01:13:05.59Z" } wheels = [ @@ -121,6 +122,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ed/20/bc79bc575ba2e2a7f70e8a1155618bb1301eaa5132a8271373a6903f73f8/babel-2.16.0-py3-none-any.whl", hash = "sha256:368b5b98b37c06b7daf6696391c3240c938b37767d4584413e8438c5c435fa8b", size = 9587599, upload-time = "2024-08-08T14:25:42.686Z" }, ] +[[package]] +name = "backports-asyncio-runner" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8e/ff/70dca7d7cb1cbc0edb2c6cc0c38b65cba36cccc491eca64cabd5fe7f8670/backports_asyncio_runner-1.2.0.tar.gz", hash = "sha256:a5aa7b2b7d8f8bfcaa2b57313f70792df84e32a2a746f585213373f900b42162", size = 69893, upload-time = "2025-07-02T02:27:15.685Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/59/76ab57e3fe74484f48a53f8e337171b4a2349e506eabe136d7e01d059086/backports_asyncio_runner-1.2.0-py3-none-any.whl", hash = "sha256:0da0a936a8aeb554eccb426dc55af3ba63bcdc69fa1a600b5bb305413a4477b5", size = 12313, upload-time = "2025-07-02T02:27:14.263Z" }, +] + [[package]] name = "beautifulsoup4" version = "4.12.3" @@ -144,59 +154,84 @@ wheels = [ [[package]] name = "cffi" -version = "1.17.1" +version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fc/97/c783634659c2920c3fc70419e3af40972dbaf758daa229a7d6ea6135c90d/cffi-1.17.1.tar.gz", hash = "sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824", size = 516621, upload-time = "2024-09-04T20:45:21.852Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/90/07/f44ca684db4e4f08a3fdc6eeb9a0d15dc6883efc7b8c90357fdbf74e186c/cffi-1.17.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:df8b1c11f177bc2313ec4b2d46baec87a5f3e71fc8b45dab2ee7cae86d9aba14", size = 182191, upload-time = "2024-09-04T20:43:30.027Z" }, - { url = "https://files.pythonhosted.org/packages/08/fd/cc2fedbd887223f9f5d170c96e57cbf655df9831a6546c1727ae13fa977a/cffi-1.17.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8f2cdc858323644ab277e9bb925ad72ae0e67f69e804f4898c070998d50b1a67", size = 178592, upload-time = "2024-09-04T20:43:32.108Z" }, - { url = "https://files.pythonhosted.org/packages/de/cc/4635c320081c78d6ffc2cab0a76025b691a91204f4aa317d568ff9280a2d/cffi-1.17.1-cp310-cp310-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:edae79245293e15384b51f88b00613ba9f7198016a5948b5dddf4917d4d26382", size = 426024, upload-time = "2024-09-04T20:43:34.186Z" }, - { url = "https://files.pythonhosted.org/packages/b6/7b/3b2b250f3aab91abe5f8a51ada1b717935fdaec53f790ad4100fe2ec64d1/cffi-1.17.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:45398b671ac6d70e67da8e4224a065cec6a93541bb7aebe1b198a61b58c7b702", size = 448188, upload-time = "2024-09-04T20:43:36.286Z" }, - { url = "https://files.pythonhosted.org/packages/d3/48/1b9283ebbf0ec065148d8de05d647a986c5f22586b18120020452fff8f5d/cffi-1.17.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ad9413ccdeda48c5afdae7e4fa2192157e991ff761e7ab8fdd8926f40b160cc3", size = 455571, upload-time = "2024-09-04T20:43:38.586Z" }, - { url = "https://files.pythonhosted.org/packages/40/87/3b8452525437b40f39ca7ff70276679772ee7e8b394934ff60e63b7b090c/cffi-1.17.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5da5719280082ac6bd9aa7becb3938dc9f9cbd57fac7d2871717b1feb0902ab6", size = 436687, upload-time = "2024-09-04T20:43:40.084Z" }, - { url = "https://files.pythonhosted.org/packages/8d/fb/4da72871d177d63649ac449aec2e8a29efe0274035880c7af59101ca2232/cffi-1.17.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2bb1a08b8008b281856e5971307cc386a8e9c5b625ac297e853d36da6efe9c17", size = 446211, upload-time = "2024-09-04T20:43:41.526Z" }, - { url = "https://files.pythonhosted.org/packages/ab/a0/62f00bcb411332106c02b663b26f3545a9ef136f80d5df746c05878f8c4b/cffi-1.17.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:045d61c734659cc045141be4bae381a41d89b741f795af1dd018bfb532fd0df8", size = 461325, upload-time = "2024-09-04T20:43:43.117Z" }, - { url = "https://files.pythonhosted.org/packages/36/83/76127035ed2e7e27b0787604d99da630ac3123bfb02d8e80c633f218a11d/cffi-1.17.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:6883e737d7d9e4899a8a695e00ec36bd4e5e4f18fabe0aca0efe0a4b44cdb13e", size = 438784, upload-time = "2024-09-04T20:43:45.256Z" }, - { url = "https://files.pythonhosted.org/packages/21/81/a6cd025db2f08ac88b901b745c163d884641909641f9b826e8cb87645942/cffi-1.17.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:6b8b4a92e1c65048ff98cfe1f735ef8f1ceb72e3d5f0c25fdb12087a23da22be", size = 461564, upload-time = "2024-09-04T20:43:46.779Z" }, - { url = "https://files.pythonhosted.org/packages/f8/fe/4d41c2f200c4a457933dbd98d3cf4e911870877bd94d9656cc0fcb390681/cffi-1.17.1-cp310-cp310-win32.whl", hash = "sha256:c9c3d058ebabb74db66e431095118094d06abf53284d9c81f27300d0e0d8bc7c", size = 171804, upload-time = "2024-09-04T20:43:48.186Z" }, - { url = "https://files.pythonhosted.org/packages/d1/b6/0b0f5ab93b0df4acc49cae758c81fe4e5ef26c3ae2e10cc69249dfd8b3ab/cffi-1.17.1-cp310-cp310-win_amd64.whl", hash = "sha256:0f048dcf80db46f0098ccac01132761580d28e28bc0f78ae0d58048063317e15", size = 181299, upload-time = "2024-09-04T20:43:49.812Z" }, - { url = "https://files.pythonhosted.org/packages/6b/f4/927e3a8899e52a27fa57a48607ff7dc91a9ebe97399b357b85a0c7892e00/cffi-1.17.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a45e3c6913c5b87b3ff120dcdc03f6131fa0065027d0ed7ee6190736a74cd401", size = 182264, upload-time = "2024-09-04T20:43:51.124Z" }, - { url = "https://files.pythonhosted.org/packages/6c/f5/6c3a8efe5f503175aaddcbea6ad0d2c96dad6f5abb205750d1b3df44ef29/cffi-1.17.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:30c5e0cb5ae493c04c8b42916e52ca38079f1b235c2f8ae5f4527b963c401caf", size = 178651, upload-time = "2024-09-04T20:43:52.872Z" }, - { url = "https://files.pythonhosted.org/packages/94/dd/a3f0118e688d1b1a57553da23b16bdade96d2f9bcda4d32e7d2838047ff7/cffi-1.17.1-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f75c7ab1f9e4aca5414ed4d8e5c0e303a34f4421f8a0d47a4d019ceff0ab6af4", size = 445259, upload-time = "2024-09-04T20:43:56.123Z" }, - { url = "https://files.pythonhosted.org/packages/2e/ea/70ce63780f096e16ce8588efe039d3c4f91deb1dc01e9c73a287939c79a6/cffi-1.17.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a1ed2dd2972641495a3ec98445e09766f077aee98a1c896dcb4ad0d303628e41", size = 469200, upload-time = "2024-09-04T20:43:57.891Z" }, - { url = "https://files.pythonhosted.org/packages/1c/a0/a4fa9f4f781bda074c3ddd57a572b060fa0df7655d2a4247bbe277200146/cffi-1.17.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:46bf43160c1a35f7ec506d254e5c890f3c03648a4dbac12d624e4490a7046cd1", size = 477235, upload-time = "2024-09-04T20:44:00.18Z" }, - { url = "https://files.pythonhosted.org/packages/62/12/ce8710b5b8affbcdd5c6e367217c242524ad17a02fe5beec3ee339f69f85/cffi-1.17.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a24ed04c8ffd54b0729c07cee15a81d964e6fee0e3d4d342a27b020d22959dc6", size = 459721, upload-time = "2024-09-04T20:44:01.585Z" }, - { url = "https://files.pythonhosted.org/packages/ff/6b/d45873c5e0242196f042d555526f92aa9e0c32355a1be1ff8c27f077fd37/cffi-1.17.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:610faea79c43e44c71e1ec53a554553fa22321b65fae24889706c0a84d4ad86d", size = 467242, upload-time = "2024-09-04T20:44:03.467Z" }, - { url = "https://files.pythonhosted.org/packages/1a/52/d9a0e523a572fbccf2955f5abe883cfa8bcc570d7faeee06336fbd50c9fc/cffi-1.17.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:a9b15d491f3ad5d692e11f6b71f7857e7835eb677955c00cc0aefcd0669adaf6", size = 477999, upload-time = "2024-09-04T20:44:05.023Z" }, - { url = "https://files.pythonhosted.org/packages/44/74/f2a2460684a1a2d00ca799ad880d54652841a780c4c97b87754f660c7603/cffi-1.17.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:de2ea4b5833625383e464549fec1bc395c1bdeeb5f25c4a3a82b5a8c756ec22f", size = 454242, upload-time = "2024-09-04T20:44:06.444Z" }, - { url = "https://files.pythonhosted.org/packages/f8/4a/34599cac7dfcd888ff54e801afe06a19c17787dfd94495ab0c8d35fe99fb/cffi-1.17.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:fc48c783f9c87e60831201f2cce7f3b2e4846bf4d8728eabe54d60700b318a0b", size = 478604, upload-time = "2024-09-04T20:44:08.206Z" }, - { url = "https://files.pythonhosted.org/packages/34/33/e1b8a1ba29025adbdcda5fb3a36f94c03d771c1b7b12f726ff7fef2ebe36/cffi-1.17.1-cp311-cp311-win32.whl", hash = "sha256:85a950a4ac9c359340d5963966e3e0a94a676bd6245a4b55bc43949eee26a655", size = 171727, upload-time = "2024-09-04T20:44:09.481Z" }, - { url = "https://files.pythonhosted.org/packages/3d/97/50228be003bb2802627d28ec0627837ac0bf35c90cf769812056f235b2d1/cffi-1.17.1-cp311-cp311-win_amd64.whl", hash = "sha256:caaf0640ef5f5517f49bc275eca1406b0ffa6aa184892812030f04c2abf589a0", size = 181400, upload-time = "2024-09-04T20:44:10.873Z" }, - { url = "https://files.pythonhosted.org/packages/5a/84/e94227139ee5fb4d600a7a4927f322e1d4aea6fdc50bd3fca8493caba23f/cffi-1.17.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:805b4371bf7197c329fcb3ead37e710d1bca9da5d583f5073b799d5c5bd1eee4", size = 183178, upload-time = "2024-09-04T20:44:12.232Z" }, - { url = "https://files.pythonhosted.org/packages/da/ee/fb72c2b48656111c4ef27f0f91da355e130a923473bf5ee75c5643d00cca/cffi-1.17.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:733e99bc2df47476e3848417c5a4540522f234dfd4ef3ab7fafdf555b082ec0c", size = 178840, upload-time = "2024-09-04T20:44:13.739Z" }, - { url = "https://files.pythonhosted.org/packages/cc/b6/db007700f67d151abadf508cbfd6a1884f57eab90b1bb985c4c8c02b0f28/cffi-1.17.1-cp312-cp312-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1257bdabf294dceb59f5e70c64a3e2f462c30c7ad68092d01bbbfb1c16b1ba36", size = 454803, upload-time = "2024-09-04T20:44:15.231Z" }, - { url = "https://files.pythonhosted.org/packages/1a/df/f8d151540d8c200eb1c6fba8cd0dfd40904f1b0682ea705c36e6c2e97ab3/cffi-1.17.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da95af8214998d77a98cc14e3a3bd00aa191526343078b530ceb0bd710fb48a5", size = 478850, upload-time = "2024-09-04T20:44:17.188Z" }, - { url = "https://files.pythonhosted.org/packages/28/c0/b31116332a547fd2677ae5b78a2ef662dfc8023d67f41b2a83f7c2aa78b1/cffi-1.17.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d63afe322132c194cf832bfec0dc69a99fb9bb6bbd550f161a49e9e855cc78ff", size = 485729, upload-time = "2024-09-04T20:44:18.688Z" }, - { url = "https://files.pythonhosted.org/packages/91/2b/9a1ddfa5c7f13cab007a2c9cc295b70fbbda7cb10a286aa6810338e60ea1/cffi-1.17.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f79fc4fc25f1c8698ff97788206bb3c2598949bfe0fef03d299eb1b5356ada99", size = 471256, upload-time = "2024-09-04T20:44:20.248Z" }, - { url = "https://files.pythonhosted.org/packages/b2/d5/da47df7004cb17e4955df6a43d14b3b4ae77737dff8bf7f8f333196717bf/cffi-1.17.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b62ce867176a75d03a665bad002af8e6d54644fad99a3c70905c543130e39d93", size = 479424, upload-time = "2024-09-04T20:44:21.673Z" }, - { url = "https://files.pythonhosted.org/packages/0b/ac/2a28bcf513e93a219c8a4e8e125534f4f6db03e3179ba1c45e949b76212c/cffi-1.17.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:386c8bf53c502fff58903061338ce4f4950cbdcb23e2902d86c0f722b786bbe3", size = 484568, upload-time = "2024-09-04T20:44:23.245Z" }, - { url = "https://files.pythonhosted.org/packages/d4/38/ca8a4f639065f14ae0f1d9751e70447a261f1a30fa7547a828ae08142465/cffi-1.17.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4ceb10419a9adf4460ea14cfd6bc43d08701f0835e979bf821052f1805850fe8", size = 488736, upload-time = "2024-09-04T20:44:24.757Z" }, - { url = "https://files.pythonhosted.org/packages/86/c5/28b2d6f799ec0bdecf44dced2ec5ed43e0eb63097b0f58c293583b406582/cffi-1.17.1-cp312-cp312-win32.whl", hash = "sha256:a08d7e755f8ed21095a310a693525137cfe756ce62d066e53f502a83dc550f65", size = 172448, upload-time = "2024-09-04T20:44:26.208Z" }, - { url = "https://files.pythonhosted.org/packages/50/b9/db34c4755a7bd1cb2d1603ac3863f22bcecbd1ba29e5ee841a4bc510b294/cffi-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:51392eae71afec0d0c8fb1a53b204dbb3bcabcb3c9b807eedf3e1e6ccf2de903", size = 181976, upload-time = "2024-09-04T20:44:27.578Z" }, - { url = "https://files.pythonhosted.org/packages/8d/f8/dd6c246b148639254dad4d6803eb6a54e8c85c6e11ec9df2cffa87571dbe/cffi-1.17.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f3a2b4222ce6b60e2e8b337bb9596923045681d71e5a082783484d845390938e", size = 182989, upload-time = "2024-09-04T20:44:28.956Z" }, - { url = "https://files.pythonhosted.org/packages/8b/f1/672d303ddf17c24fc83afd712316fda78dc6fce1cd53011b839483e1ecc8/cffi-1.17.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0984a4925a435b1da406122d4d7968dd861c1385afe3b45ba82b750f229811e2", size = 178802, upload-time = "2024-09-04T20:44:30.289Z" }, - { url = "https://files.pythonhosted.org/packages/0e/2d/eab2e858a91fdff70533cab61dcff4a1f55ec60425832ddfdc9cd36bc8af/cffi-1.17.1-cp313-cp313-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d01b12eeeb4427d3110de311e1774046ad344f5b1a7403101878976ecd7a10f3", size = 454792, upload-time = "2024-09-04T20:44:32.01Z" }, - { url = "https://files.pythonhosted.org/packages/75/b2/fbaec7c4455c604e29388d55599b99ebcc250a60050610fadde58932b7ee/cffi-1.17.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:706510fe141c86a69c8ddc029c7910003a17353970cff3b904ff0686a5927683", size = 478893, upload-time = "2024-09-04T20:44:33.606Z" }, - { url = "https://files.pythonhosted.org/packages/4f/b7/6e4a2162178bf1935c336d4da8a9352cccab4d3a5d7914065490f08c0690/cffi-1.17.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:de55b766c7aa2e2a3092c51e0483d700341182f08e67c63630d5b6f200bb28e5", size = 485810, upload-time = "2024-09-04T20:44:35.191Z" }, - { url = "https://files.pythonhosted.org/packages/c7/8a/1d0e4a9c26e54746dc08c2c6c037889124d4f59dffd853a659fa545f1b40/cffi-1.17.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c59d6e989d07460165cc5ad3c61f9fd8f1b4796eacbd81cee78957842b834af4", size = 471200, upload-time = "2024-09-04T20:44:36.743Z" }, - { url = "https://files.pythonhosted.org/packages/26/9f/1aab65a6c0db35f43c4d1b4f580e8df53914310afc10ae0397d29d697af4/cffi-1.17.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd398dbc6773384a17fe0d3e7eeb8d1a21c2200473ee6806bb5e6a8e62bb73dd", size = 479447, upload-time = "2024-09-04T20:44:38.492Z" }, - { url = "https://files.pythonhosted.org/packages/5f/e4/fb8b3dd8dc0e98edf1135ff067ae070bb32ef9d509d6cb0f538cd6f7483f/cffi-1.17.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3edc8d958eb099c634dace3c7e16560ae474aa3803a5df240542b305d14e14ed", size = 484358, upload-time = "2024-09-04T20:44:40.046Z" }, - { url = "https://files.pythonhosted.org/packages/f1/47/d7145bf2dc04684935d57d67dff9d6d795b2ba2796806bb109864be3a151/cffi-1.17.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:72e72408cad3d5419375fc87d289076ee319835bdfa2caad331e377589aebba9", size = 488469, upload-time = "2024-09-04T20:44:41.616Z" }, - { url = "https://files.pythonhosted.org/packages/bf/ee/f94057fa6426481d663b88637a9a10e859e492c73d0384514a17d78ee205/cffi-1.17.1-cp313-cp313-win32.whl", hash = "sha256:e03eab0a8677fa80d646b5ddece1cbeaf556c313dcfac435ba11f107ba117b5d", size = 172475, upload-time = "2024-09-04T20:44:43.733Z" }, - { url = "https://files.pythonhosted.org/packages/7c/fc/6a8cb64e5f0324877d503c854da15d76c1e50eb722e320b15345c4d0c6de/cffi-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:f6a16c31041f09ead72d69f583767292f750d24913dadacf5756b966aacb3f1a", size = 182009, upload-time = "2024-09-04T20:44:45.309Z" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/d7/516d984057745a6cd96575eea814fe1edd6646ee6efd552fb7b0921dec83/cffi-2.0.0-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:0cf2d91ecc3fcc0625c2c530fe004f82c110405f101548512cce44322fa8ac44", size = 184283, upload-time = "2025-09-08T23:22:08.01Z" }, + { url = "https://files.pythonhosted.org/packages/9e/84/ad6a0b408daa859246f57c03efd28e5dd1b33c21737c2db84cae8c237aa5/cffi-2.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:f73b96c41e3b2adedc34a7356e64c8eb96e03a3782b535e043a986276ce12a49", size = 180504, upload-time = "2025-09-08T23:22:10.637Z" }, + { url = "https://files.pythonhosted.org/packages/50/bd/b1a6362b80628111e6653c961f987faa55262b4002fcec42308cad1db680/cffi-2.0.0-cp310-cp310-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:53f77cbe57044e88bbd5ed26ac1d0514d2acf0591dd6bb02a3ae37f76811b80c", size = 208811, upload-time = "2025-09-08T23:22:12.267Z" }, + { url = "https://files.pythonhosted.org/packages/4f/27/6933a8b2562d7bd1fb595074cf99cc81fc3789f6a6c05cdabb46284a3188/cffi-2.0.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3e837e369566884707ddaf85fc1744b47575005c0a229de3327f8f9a20f4efeb", size = 216402, upload-time = "2025-09-08T23:22:13.455Z" }, + { url = "https://files.pythonhosted.org/packages/05/eb/b86f2a2645b62adcfff53b0dd97e8dfafb5c8aa864bd0d9a2c2049a0d551/cffi-2.0.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:5eda85d6d1879e692d546a078b44251cdd08dd1cfb98dfb77b670c97cee49ea0", size = 203217, upload-time = "2025-09-08T23:22:14.596Z" }, + { url = "https://files.pythonhosted.org/packages/9f/e0/6cbe77a53acf5acc7c08cc186c9928864bd7c005f9efd0d126884858a5fe/cffi-2.0.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9332088d75dc3241c702d852d4671613136d90fa6881da7d770a483fd05248b4", size = 203079, upload-time = "2025-09-08T23:22:15.769Z" }, + { url = "https://files.pythonhosted.org/packages/98/29/9b366e70e243eb3d14a5cb488dfd3a0b6b2f1fb001a203f653b93ccfac88/cffi-2.0.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc7de24befaeae77ba923797c7c87834c73648a05a4bde34b3b7e5588973a453", size = 216475, upload-time = "2025-09-08T23:22:17.427Z" }, + { url = "https://files.pythonhosted.org/packages/21/7a/13b24e70d2f90a322f2900c5d8e1f14fa7e2a6b3332b7309ba7b2ba51a5a/cffi-2.0.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cf364028c016c03078a23b503f02058f1814320a56ad535686f90565636a9495", size = 218829, upload-time = "2025-09-08T23:22:19.069Z" }, + { url = "https://files.pythonhosted.org/packages/60/99/c9dc110974c59cc981b1f5b66e1d8af8af764e00f0293266824d9c4254bc/cffi-2.0.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:e11e82b744887154b182fd3e7e8512418446501191994dbf9c9fc1f32cc8efd5", size = 211211, upload-time = "2025-09-08T23:22:20.588Z" }, + { url = "https://files.pythonhosted.org/packages/49/72/ff2d12dbf21aca1b32a40ed792ee6b40f6dc3a9cf1644bd7ef6e95e0ac5e/cffi-2.0.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8ea985900c5c95ce9db1745f7933eeef5d314f0565b27625d9a10ec9881e1bfb", size = 218036, upload-time = "2025-09-08T23:22:22.143Z" }, + { url = "https://files.pythonhosted.org/packages/e2/cc/027d7fb82e58c48ea717149b03bcadcbdc293553edb283af792bd4bcbb3f/cffi-2.0.0-cp310-cp310-win32.whl", hash = "sha256:1f72fb8906754ac8a2cc3f9f5aaa298070652a0ffae577e0ea9bd480dc3c931a", size = 172184, upload-time = "2025-09-08T23:22:23.328Z" }, + { url = "https://files.pythonhosted.org/packages/33/fa/072dd15ae27fbb4e06b437eb6e944e75b068deb09e2a2826039e49ee2045/cffi-2.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:b18a3ed7d5b3bd8d9ef7a8cb226502c6bf8308df1525e1cc676c3680e7176739", size = 182790, upload-time = "2025-09-08T23:22:24.752Z" }, + { url = "https://files.pythonhosted.org/packages/12/4a/3dfd5f7850cbf0d06dc84ba9aa00db766b52ca38d8b86e3a38314d52498c/cffi-2.0.0-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:b4c854ef3adc177950a8dfc81a86f5115d2abd545751a304c5bcf2c2c7283cfe", size = 184344, upload-time = "2025-09-08T23:22:26.456Z" }, + { url = "https://files.pythonhosted.org/packages/4f/8b/f0e4c441227ba756aafbe78f117485b25bb26b1c059d01f137fa6d14896b/cffi-2.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2de9a304e27f7596cd03d16f1b7c72219bd944e99cc52b84d0145aefb07cbd3c", size = 180560, upload-time = "2025-09-08T23:22:28.197Z" }, + { url = "https://files.pythonhosted.org/packages/b1/b7/1200d354378ef52ec227395d95c2576330fd22a869f7a70e88e1447eb234/cffi-2.0.0-cp311-cp311-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:baf5215e0ab74c16e2dd324e8ec067ef59e41125d3eade2b863d294fd5035c92", size = 209613, upload-time = "2025-09-08T23:22:29.475Z" }, + { url = "https://files.pythonhosted.org/packages/b8/56/6033f5e86e8cc9bb629f0077ba71679508bdf54a9a5e112a3c0b91870332/cffi-2.0.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:730cacb21e1bdff3ce90babf007d0a0917cc3e6492f336c2f0134101e0944f93", size = 216476, upload-time = "2025-09-08T23:22:31.063Z" }, + { url = "https://files.pythonhosted.org/packages/dc/7f/55fecd70f7ece178db2f26128ec41430d8720f2d12ca97bf8f0a628207d5/cffi-2.0.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:6824f87845e3396029f3820c206e459ccc91760e8fa24422f8b0c3d1731cbec5", size = 203374, upload-time = "2025-09-08T23:22:32.507Z" }, + { url = "https://files.pythonhosted.org/packages/84/ef/a7b77c8bdc0f77adc3b46888f1ad54be8f3b7821697a7b89126e829e676a/cffi-2.0.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9de40a7b0323d889cf8d23d1ef214f565ab154443c42737dfe52ff82cf857664", size = 202597, upload-time = "2025-09-08T23:22:34.132Z" }, + { url = "https://files.pythonhosted.org/packages/d7/91/500d892b2bf36529a75b77958edfcd5ad8e2ce4064ce2ecfeab2125d72d1/cffi-2.0.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8941aaadaf67246224cee8c3803777eed332a19d909b47e29c9842ef1e79ac26", size = 215574, upload-time = "2025-09-08T23:22:35.443Z" }, + { url = "https://files.pythonhosted.org/packages/44/64/58f6255b62b101093d5df22dcb752596066c7e89dd725e0afaed242a61be/cffi-2.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:a05d0c237b3349096d3981b727493e22147f934b20f6f125a3eba8f994bec4a9", size = 218971, upload-time = "2025-09-08T23:22:36.805Z" }, + { url = "https://files.pythonhosted.org/packages/ab/49/fa72cebe2fd8a55fbe14956f9970fe8eb1ac59e5df042f603ef7c8ba0adc/cffi-2.0.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:94698a9c5f91f9d138526b48fe26a199609544591f859c870d477351dc7b2414", size = 211972, upload-time = "2025-09-08T23:22:38.436Z" }, + { url = "https://files.pythonhosted.org/packages/0b/28/dd0967a76aab36731b6ebfe64dec4e981aff7e0608f60c2d46b46982607d/cffi-2.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:5fed36fccc0612a53f1d4d9a816b50a36702c28a2aa880cb8a122b3466638743", size = 217078, upload-time = "2025-09-08T23:22:39.776Z" }, + { url = "https://files.pythonhosted.org/packages/2b/c0/015b25184413d7ab0a410775fdb4a50fca20f5589b5dab1dbbfa3baad8ce/cffi-2.0.0-cp311-cp311-win32.whl", hash = "sha256:c649e3a33450ec82378822b3dad03cc228b8f5963c0c12fc3b1e0ab940f768a5", size = 172076, upload-time = "2025-09-08T23:22:40.95Z" }, + { url = "https://files.pythonhosted.org/packages/ae/8f/dc5531155e7070361eb1b7e4c1a9d896d0cb21c49f807a6c03fd63fc877e/cffi-2.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:66f011380d0e49ed280c789fbd08ff0d40968ee7b665575489afa95c98196ab5", size = 182820, upload-time = "2025-09-08T23:22:42.463Z" }, + { url = "https://files.pythonhosted.org/packages/95/5c/1b493356429f9aecfd56bc171285a4c4ac8697f76e9bbbbb105e537853a1/cffi-2.0.0-cp311-cp311-win_arm64.whl", hash = "sha256:c6638687455baf640e37344fe26d37c404db8b80d037c3d29f58fe8d1c3b194d", size = 177635, upload-time = "2025-09-08T23:22:43.623Z" }, + { url = "https://files.pythonhosted.org/packages/ea/47/4f61023ea636104d4f16ab488e268b93008c3d0bb76893b1b31db1f96802/cffi-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d02d6655b0e54f54c4ef0b94eb6be0607b70853c45ce98bd278dc7de718be5d", size = 185271, upload-time = "2025-09-08T23:22:44.795Z" }, + { url = "https://files.pythonhosted.org/packages/df/a2/781b623f57358e360d62cdd7a8c681f074a71d445418a776eef0aadb4ab4/cffi-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8eca2a813c1cb7ad4fb74d368c2ffbbb4789d377ee5bb8df98373c2cc0dee76c", size = 181048, upload-time = "2025-09-08T23:22:45.938Z" }, + { url = "https://files.pythonhosted.org/packages/ff/df/a4f0fbd47331ceeba3d37c2e51e9dfc9722498becbeec2bd8bc856c9538a/cffi-2.0.0-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:21d1152871b019407d8ac3985f6775c079416c282e431a4da6afe7aefd2bccbe", size = 212529, upload-time = "2025-09-08T23:22:47.349Z" }, + { url = "https://files.pythonhosted.org/packages/d5/72/12b5f8d3865bf0f87cf1404d8c374e7487dcf097a1c91c436e72e6badd83/cffi-2.0.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b21e08af67b8a103c71a250401c78d5e0893beff75e28c53c98f4de42f774062", size = 220097, upload-time = "2025-09-08T23:22:48.677Z" }, + { url = "https://files.pythonhosted.org/packages/c2/95/7a135d52a50dfa7c882ab0ac17e8dc11cec9d55d2c18dda414c051c5e69e/cffi-2.0.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:1e3a615586f05fc4065a8b22b8152f0c1b00cdbc60596d187c2a74f9e3036e4e", size = 207983, upload-time = "2025-09-08T23:22:50.06Z" }, + { url = "https://files.pythonhosted.org/packages/3a/c8/15cb9ada8895957ea171c62dc78ff3e99159ee7adb13c0123c001a2546c1/cffi-2.0.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:81afed14892743bbe14dacb9e36d9e0e504cd204e0b165062c488942b9718037", size = 206519, upload-time = "2025-09-08T23:22:51.364Z" }, + { url = "https://files.pythonhosted.org/packages/78/2d/7fa73dfa841b5ac06c7b8855cfc18622132e365f5b81d02230333ff26e9e/cffi-2.0.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3e17ed538242334bf70832644a32a7aae3d83b57567f9fd60a26257e992b79ba", size = 219572, upload-time = "2025-09-08T23:22:52.902Z" }, + { url = "https://files.pythonhosted.org/packages/07/e0/267e57e387b4ca276b90f0434ff88b2c2241ad72b16d31836adddfd6031b/cffi-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3925dd22fa2b7699ed2617149842d2e6adde22b262fcbfada50e3d195e4b3a94", size = 222963, upload-time = "2025-09-08T23:22:54.518Z" }, + { url = "https://files.pythonhosted.org/packages/b6/75/1f2747525e06f53efbd878f4d03bac5b859cbc11c633d0fb81432d98a795/cffi-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c8f814d84194c9ea681642fd164267891702542f028a15fc97d4674b6206187", size = 221361, upload-time = "2025-09-08T23:22:55.867Z" }, + { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" }, + { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" }, + { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" }, + { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" }, + { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" }, + { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" }, + { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" }, + { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" }, + { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" }, + { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" }, + { url = "https://files.pythonhosted.org/packages/eb/6d/bf9bda840d5f1dfdbf0feca87fbdb64a918a69bca42cfa0ba7b137c48cb8/cffi-2.0.0-cp313-cp313-win32.whl", hash = "sha256:74a03b9698e198d47562765773b4a8309919089150a0bb17d829ad7b44b60d27", size = 172909, upload-time = "2025-09-08T23:23:14.32Z" }, + { url = "https://files.pythonhosted.org/packages/37/18/6519e1ee6f5a1e579e04b9ddb6f1676c17368a7aba48299c3759bbc3c8b3/cffi-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:19f705ada2530c1167abacb171925dd886168931e0a7b78f5bffcae5c6b5be75", size = 183402, upload-time = "2025-09-08T23:23:15.535Z" }, + { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" }, + { url = "https://files.pythonhosted.org/packages/92/c4/3ce07396253a83250ee98564f8d7e9789fab8e58858f35d07a9a2c78de9f/cffi-2.0.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fc33c5141b55ed366cfaad382df24fe7dcbc686de5be719b207bb248e3053dc5", size = 185320, upload-time = "2025-09-08T23:23:18.087Z" }, + { url = "https://files.pythonhosted.org/packages/59/dd/27e9fa567a23931c838c6b02d0764611c62290062a6d4e8ff7863daf9730/cffi-2.0.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c654de545946e0db659b3400168c9ad31b5d29593291482c43e3564effbcee13", size = 181487, upload-time = "2025-09-08T23:23:19.622Z" }, + { url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" }, + { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" }, + { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" }, + { url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" }, + { url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" }, + { url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" }, + { url = "https://files.pythonhosted.org/packages/3e/aa/df335faa45b395396fcbc03de2dfcab242cd61a9900e914fe682a59170b1/cffi-2.0.0-cp314-cp314-win32.whl", hash = "sha256:087067fa8953339c723661eda6b54bc98c5625757ea62e95eb4898ad5e776e9f", size = 175328, upload-time = "2025-09-08T23:23:44.61Z" }, + { url = "https://files.pythonhosted.org/packages/bb/92/882c2d30831744296ce713f0feb4c1cd30f346ef747b530b5318715cc367/cffi-2.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:203a48d1fb583fc7d78a4c6655692963b860a417c0528492a6bc21f1aaefab25", size = 185650, upload-time = "2025-09-08T23:23:45.848Z" }, + { url = "https://files.pythonhosted.org/packages/9f/2c/98ece204b9d35a7366b5b2c6539c350313ca13932143e79dc133ba757104/cffi-2.0.0-cp314-cp314-win_arm64.whl", hash = "sha256:dbd5c7a25a7cb98f5ca55d258b103a2054f859a46ae11aaf23134f9cc0d356ad", size = 180687, upload-time = "2025-09-08T23:23:47.105Z" }, + { url = "https://files.pythonhosted.org/packages/3e/61/c768e4d548bfa607abcda77423448df8c471f25dbe64fb2ef6d555eae006/cffi-2.0.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:9a67fc9e8eb39039280526379fb3a70023d77caec1852002b4da7e8b270c4dd9", size = 188773, upload-time = "2025-09-08T23:23:29.347Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ea/5f76bce7cf6fcd0ab1a1058b5af899bfbef198bea4d5686da88471ea0336/cffi-2.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7a66c7204d8869299919db4d5069a82f1561581af12b11b3c9f48c584eb8743d", size = 185013, upload-time = "2025-09-08T23:23:30.63Z" }, + { url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" }, + { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" }, + { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" }, + { url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" }, + { url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" }, + { url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" }, + { url = "https://files.pythonhosted.org/packages/a0/1d/ec1a60bd1a10daa292d3cd6bb0b359a81607154fb8165f3ec95fe003b85c/cffi-2.0.0-cp314-cp314t-win32.whl", hash = "sha256:1fc9ea04857caf665289b7a75923f2c6ed559b8298a1b8c49e59f7dd95c8481e", size = 180487, upload-time = "2025-09-08T23:23:40.423Z" }, + { url = "https://files.pythonhosted.org/packages/bf/41/4c1168c74fac325c0c8156f04b6749c8b6a8f405bbf91413ba088359f60d/cffi-2.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:d68b6cef7827e8641e8ef16f4494edda8b36104d79773a334beaa1e3521430f6", size = 191726, upload-time = "2025-09-08T23:23:41.742Z" }, + { url = "https://files.pythonhosted.org/packages/ae/3a/dbeec9d1ee0844c679f6bb5d6ad4e9f198b1224f4e7a32825f47f6192b0c/cffi-2.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0a1527a803f0a659de1af2e1fd700213caba79377e27e4693648c2923da066f9", size = 184195, upload-time = "2025-09-08T23:23:43.004Z" }, ] [[package]] @@ -298,39 +333,62 @@ wheels = [ [[package]] name = "cryptography" -version = "44.0.0" +version = "46.0.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/91/4c/45dfa6829acffa344e3967d6006ee4ae8be57af746ae2eba1c431949b32c/cryptography-44.0.0.tar.gz", hash = "sha256:cd4e834f340b4293430701e772ec543b0fbe6c2dea510a5286fe0acabe153a02", size = 710657, upload-time = "2024-11-27T18:07:10.168Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/55/09/8cc67f9b84730ad330b3b72cf867150744bf07ff113cda21a15a1c6d2c7c/cryptography-44.0.0-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:84111ad4ff3f6253820e6d3e58be2cc2a00adb29335d4cacb5ab4d4d34f2a123", size = 6541833, upload-time = "2024-11-27T18:05:55.475Z" }, - { url = "https://files.pythonhosted.org/packages/7e/5b/3759e30a103144e29632e7cb72aec28cedc79e514b2ea8896bb17163c19b/cryptography-44.0.0-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b15492a11f9e1b62ba9d73c210e2416724633167de94607ec6069ef724fad092", size = 3922710, upload-time = "2024-11-27T18:05:58.621Z" }, - { url = "https://files.pythonhosted.org/packages/5f/58/3b14bf39f1a0cfd679e753e8647ada56cddbf5acebffe7db90e184c76168/cryptography-44.0.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:831c3c4d0774e488fdc83a1923b49b9957d33287de923d58ebd3cec47a0ae43f", size = 4137546, upload-time = "2024-11-27T18:06:01.062Z" }, - { url = "https://files.pythonhosted.org/packages/98/65/13d9e76ca19b0ba5603d71ac8424b5694415b348e719db277b5edc985ff5/cryptography-44.0.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:761817a3377ef15ac23cd7834715081791d4ec77f9297ee694ca1ee9c2c7e5eb", size = 3915420, upload-time = "2024-11-27T18:06:03.487Z" }, - { url = "https://files.pythonhosted.org/packages/b1/07/40fe09ce96b91fc9276a9ad272832ead0fddedcba87f1190372af8e3039c/cryptography-44.0.0-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3c672a53c0fb4725a29c303be906d3c1fa99c32f58abe008a82705f9ee96f40b", size = 4154498, upload-time = "2024-11-27T18:06:05.763Z" }, - { url = "https://files.pythonhosted.org/packages/75/ea/af65619c800ec0a7e4034207aec543acdf248d9bffba0533342d1bd435e1/cryptography-44.0.0-cp37-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:4ac4c9f37eba52cb6fbeaf5b59c152ea976726b865bd4cf87883a7e7006cc543", size = 3932569, upload-time = "2024-11-27T18:06:07.489Z" }, - { url = "https://files.pythonhosted.org/packages/c7/af/d1deb0c04d59612e3d5e54203159e284d3e7a6921e565bb0eeb6269bdd8a/cryptography-44.0.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ed3534eb1090483c96178fcb0f8893719d96d5274dfde98aa6add34614e97c8e", size = 4016721, upload-time = "2024-11-27T18:06:11.57Z" }, - { url = "https://files.pythonhosted.org/packages/bd/69/7ca326c55698d0688db867795134bdfac87136b80ef373aaa42b225d6dd5/cryptography-44.0.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f3f6fdfa89ee2d9d496e2c087cebef9d4fcbb0ad63c40e821b39f74bf48d9c5e", size = 4240915, upload-time = "2024-11-27T18:06:13.515Z" }, - { url = "https://files.pythonhosted.org/packages/ef/d4/cae11bf68c0f981e0413906c6dd03ae7fa864347ed5fac40021df1ef467c/cryptography-44.0.0-cp37-abi3-win32.whl", hash = "sha256:eb33480f1bad5b78233b0ad3e1b0be21e8ef1da745d8d2aecbb20671658b9053", size = 2757925, upload-time = "2024-11-27T18:06:16.019Z" }, - { url = "https://files.pythonhosted.org/packages/64/b1/50d7739254d2002acae64eed4fc43b24ac0cc44bf0a0d388d1ca06ec5bb1/cryptography-44.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:abc998e0c0eee3c8a1904221d3f67dcfa76422b23620173e28c11d3e626c21bd", size = 3202055, upload-time = "2024-11-27T18:06:19.113Z" }, - { url = "https://files.pythonhosted.org/packages/11/18/61e52a3d28fc1514a43b0ac291177acd1b4de00e9301aaf7ef867076ff8a/cryptography-44.0.0-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:660cb7312a08bc38be15b696462fa7cc7cd85c3ed9c576e81f4dc4d8b2b31591", size = 6542801, upload-time = "2024-11-27T18:06:21.431Z" }, - { url = "https://files.pythonhosted.org/packages/1a/07/5f165b6c65696ef75601b781a280fc3b33f1e0cd6aa5a92d9fb96c410e97/cryptography-44.0.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1923cb251c04be85eec9fda837661c67c1049063305d6be5721643c22dd4e2b7", size = 3922613, upload-time = "2024-11-27T18:06:24.314Z" }, - { url = "https://files.pythonhosted.org/packages/28/34/6b3ac1d80fc174812486561cf25194338151780f27e438526f9c64e16869/cryptography-44.0.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:404fdc66ee5f83a1388be54300ae978b2efd538018de18556dde92575e05defc", size = 4137925, upload-time = "2024-11-27T18:06:27.079Z" }, - { url = "https://files.pythonhosted.org/packages/d0/c7/c656eb08fd22255d21bc3129625ed9cd5ee305f33752ef2278711b3fa98b/cryptography-44.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:c5eb858beed7835e5ad1faba59e865109f3e52b3783b9ac21e7e47dc5554e289", size = 3915417, upload-time = "2024-11-27T18:06:28.959Z" }, - { url = "https://files.pythonhosted.org/packages/ef/82/72403624f197af0db6bac4e58153bc9ac0e6020e57234115db9596eee85d/cryptography-44.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:f53c2c87e0fb4b0c00fa9571082a057e37690a8f12233306161c8f4b819960b7", size = 4155160, upload-time = "2024-11-27T18:06:30.866Z" }, - { url = "https://files.pythonhosted.org/packages/a2/cd/2f3c440913d4329ade49b146d74f2e9766422e1732613f57097fea61f344/cryptography-44.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:9e6fc8a08e116fb7c7dd1f040074c9d7b51d74a8ea40d4df2fc7aa08b76b9e6c", size = 3932331, upload-time = "2024-11-27T18:06:33.432Z" }, - { url = "https://files.pythonhosted.org/packages/7f/df/8be88797f0a1cca6e255189a57bb49237402b1880d6e8721690c5603ac23/cryptography-44.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:d2436114e46b36d00f8b72ff57e598978b37399d2786fd39793c36c6d5cb1c64", size = 4017372, upload-time = "2024-11-27T18:06:38.343Z" }, - { url = "https://files.pythonhosted.org/packages/af/36/5ccc376f025a834e72b8e52e18746b927f34e4520487098e283a719c205e/cryptography-44.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a01956ddfa0a6790d594f5b34fc1bfa6098aca434696a03cfdbe469b8ed79285", size = 4239657, upload-time = "2024-11-27T18:06:41.045Z" }, - { url = "https://files.pythonhosted.org/packages/46/b0/f4f7d0d0bcfbc8dd6296c1449be326d04217c57afb8b2594f017eed95533/cryptography-44.0.0-cp39-abi3-win32.whl", hash = "sha256:eca27345e1214d1b9f9490d200f9db5a874479be914199194e746c893788d417", size = 2758672, upload-time = "2024-11-27T18:06:43.566Z" }, - { url = "https://files.pythonhosted.org/packages/97/9b/443270b9210f13f6ef240eff73fd32e02d381e7103969dc66ce8e89ee901/cryptography-44.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:708ee5f1bafe76d041b53a4f95eb28cdeb8d18da17e597d46d7833ee59b97ede", size = 3202071, upload-time = "2024-11-27T18:06:45.586Z" }, - { url = "https://files.pythonhosted.org/packages/77/d4/fea74422326388bbac0c37b7489a0fcb1681a698c3b875959430ba550daa/cryptography-44.0.0-pp310-pypy310_pp73-macosx_10_9_x86_64.whl", hash = "sha256:37d76e6863da3774cd9db5b409a9ecfd2c71c981c38788d3fcfaf177f447b731", size = 3338857, upload-time = "2024-11-27T18:06:48.88Z" }, - { url = "https://files.pythonhosted.org/packages/1a/aa/ba8a7467c206cb7b62f09b4168da541b5109838627f582843bbbe0235e8e/cryptography-44.0.0-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:f677e1268c4e23420c3acade68fac427fffcb8d19d7df95ed7ad17cdef8404f4", size = 3850615, upload-time = "2024-11-27T18:06:50.774Z" }, - { url = "https://files.pythonhosted.org/packages/89/fa/b160e10a64cc395d090105be14f399b94e617c879efd401188ce0fea39ee/cryptography-44.0.0-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:f5e7cb1e5e56ca0933b4873c0220a78b773b24d40d186b6738080b73d3d0a756", size = 4081622, upload-time = "2024-11-27T18:06:55.126Z" }, - { url = "https://files.pythonhosted.org/packages/47/8f/20ff0656bb0cf7af26ec1d01f780c5cfbaa7666736063378c5f48558b515/cryptography-44.0.0-pp310-pypy310_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:8b3e6eae66cf54701ee7d9c83c30ac0a1e3fa17be486033000f2a73a12ab507c", size = 3867546, upload-time = "2024-11-27T18:06:57.694Z" }, - { url = "https://files.pythonhosted.org/packages/38/d9/28edf32ee2fcdca587146bcde90102a7319b2f2c690edfa627e46d586050/cryptography-44.0.0-pp310-pypy310_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:be4ce505894d15d5c5037167ffb7f0ae90b7be6f2a98f9a5c3442395501c32fa", size = 4090937, upload-time = "2024-11-27T18:07:00.338Z" }, - { url = "https://files.pythonhosted.org/packages/cc/9d/37e5da7519de7b0b070a3fedd4230fe76d50d2a21403e0f2153d70ac4163/cryptography-44.0.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:62901fb618f74d7d81bf408c8719e9ec14d863086efe4185afd07c352aee1d2c", size = 3128774, upload-time = "2024-11-27T18:07:02.157Z" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/47/93/ac8f3d5ff04d54bc814e961a43ae5b0b146154c89c61b47bb07557679b18/cryptography-46.0.7.tar.gz", hash = "sha256:e4cfd68c5f3e0bfdad0d38e023239b96a2fe84146481852dffbcca442c245aa5", size = 750652, upload-time = "2026-04-08T01:57:54.692Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/5d/4a8f770695d73be252331e60e526291e3df0c9b27556a90a6b47bccca4c2/cryptography-46.0.7-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:ea42cbe97209df307fdc3b155f1b6fa2577c0defa8f1f7d3be7d31d189108ad4", size = 7179869, upload-time = "2026-04-08T01:56:17.157Z" }, + { url = "https://files.pythonhosted.org/packages/5f/45/6d80dc379b0bbc1f9d1e429f42e4cb9e1d319c7a8201beffd967c516ea01/cryptography-46.0.7-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b36a4695e29fe69215d75960b22577197aca3f7a25b9cf9d165dcfe9d80bc325", size = 4275492, upload-time = "2026-04-08T01:56:19.36Z" }, + { url = "https://files.pythonhosted.org/packages/4a/9a/1765afe9f572e239c3469f2cb429f3ba7b31878c893b246b4b2994ffe2fe/cryptography-46.0.7-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5ad9ef796328c5e3c4ceed237a183f5d41d21150f972455a9d926593a1dcb308", size = 4426670, upload-time = "2026-04-08T01:56:21.415Z" }, + { url = "https://files.pythonhosted.org/packages/8f/3e/af9246aaf23cd4ee060699adab1e47ced3f5f7e7a8ffdd339f817b446462/cryptography-46.0.7-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:73510b83623e080a2c35c62c15298096e2a5dc8d51c3b4e1740211839d0dea77", size = 4280275, upload-time = "2026-04-08T01:56:23.539Z" }, + { url = "https://files.pythonhosted.org/packages/0f/54/6bbbfc5efe86f9d71041827b793c24811a017c6ac0fd12883e4caa86b8ed/cryptography-46.0.7-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:cbd5fb06b62bd0721e1170273d3f4d5a277044c47ca27ee257025146c34cbdd1", size = 4928402, upload-time = "2026-04-08T01:56:25.624Z" }, + { url = "https://files.pythonhosted.org/packages/2d/cf/054b9d8220f81509939599c8bdbc0c408dbd2bdd41688616a20731371fe0/cryptography-46.0.7-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:420b1e4109cc95f0e5700eed79908cef9268265c773d3a66f7af1eef53d409ef", size = 4459985, upload-time = "2026-04-08T01:56:27.309Z" }, + { url = "https://files.pythonhosted.org/packages/f9/46/4e4e9c6040fb01c7467d47217d2f882daddeb8828f7df800cb806d8a2288/cryptography-46.0.7-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:24402210aa54baae71d99441d15bb5a1919c195398a87b563df84468160a65de", size = 3990652, upload-time = "2026-04-08T01:56:29.095Z" }, + { url = "https://files.pythonhosted.org/packages/36/5f/313586c3be5a2fbe87e4c9a254207b860155a8e1f3cca99f9910008e7d08/cryptography-46.0.7-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:8a469028a86f12eb7d2fe97162d0634026d92a21f3ae0ac87ed1c4a447886c83", size = 4279805, upload-time = "2026-04-08T01:56:30.928Z" }, + { url = "https://files.pythonhosted.org/packages/69/33/60dfc4595f334a2082749673386a4d05e4f0cf4df8248e63b2c3437585f2/cryptography-46.0.7-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:9694078c5d44c157ef3162e3bf3946510b857df5a3955458381d1c7cfc143ddb", size = 4892883, upload-time = "2026-04-08T01:56:32.614Z" }, + { url = "https://files.pythonhosted.org/packages/c7/0b/333ddab4270c4f5b972f980adef4faa66951a4aaf646ca067af597f15563/cryptography-46.0.7-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:42a1e5f98abb6391717978baf9f90dc28a743b7d9be7f0751a6f56a75d14065b", size = 4459756, upload-time = "2026-04-08T01:56:34.306Z" }, + { url = "https://files.pythonhosted.org/packages/d2/14/633913398b43b75f1234834170947957c6b623d1701ffc7a9600da907e89/cryptography-46.0.7-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:91bbcb08347344f810cbe49065914fe048949648f6bd5c2519f34619142bbe85", size = 4410244, upload-time = "2026-04-08T01:56:35.977Z" }, + { url = "https://files.pythonhosted.org/packages/10/f2/19ceb3b3dc14009373432af0c13f46aa08e3ce334ec6eff13492e1812ccd/cryptography-46.0.7-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:5d1c02a14ceb9148cc7816249f64f623fbfee39e8c03b3650d842ad3f34d637e", size = 4674868, upload-time = "2026-04-08T01:56:38.034Z" }, + { url = "https://files.pythonhosted.org/packages/1a/bb/a5c213c19ee94b15dfccc48f363738633a493812687f5567addbcbba9f6f/cryptography-46.0.7-cp311-abi3-win32.whl", hash = "sha256:d23c8ca48e44ee015cd0a54aeccdf9f09004eba9fc96f38c911011d9ff1bd457", size = 3026504, upload-time = "2026-04-08T01:56:39.666Z" }, + { url = "https://files.pythonhosted.org/packages/2b/02/7788f9fefa1d060ca68717c3901ae7fffa21ee087a90b7f23c7a603c32ae/cryptography-46.0.7-cp311-abi3-win_amd64.whl", hash = "sha256:397655da831414d165029da9bc483bed2fe0e75dde6a1523ec2fe63f3c46046b", size = 3488363, upload-time = "2026-04-08T01:56:41.893Z" }, + { url = "https://files.pythonhosted.org/packages/7b/56/15619b210e689c5403bb0540e4cb7dbf11a6bf42e483b7644e471a2812b3/cryptography-46.0.7-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:d151173275e1728cf7839aaa80c34fe550c04ddb27b34f48c232193df8db5842", size = 7119671, upload-time = "2026-04-08T01:56:44Z" }, + { url = "https://files.pythonhosted.org/packages/74/66/e3ce040721b0b5599e175ba91ab08884c75928fbeb74597dd10ef13505d2/cryptography-46.0.7-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:db0f493b9181c7820c8134437eb8b0b4792085d37dbb24da050476ccb664e59c", size = 4268551, upload-time = "2026-04-08T01:56:46.071Z" }, + { url = "https://files.pythonhosted.org/packages/03/11/5e395f961d6868269835dee1bafec6a1ac176505a167f68b7d8818431068/cryptography-46.0.7-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ebd6daf519b9f189f85c479427bbd6e9c9037862cf8fe89ee35503bd209ed902", size = 4408887, upload-time = "2026-04-08T01:56:47.718Z" }, + { url = "https://files.pythonhosted.org/packages/40/53/8ed1cf4c3b9c8e611e7122fb56f1c32d09e1fff0f1d77e78d9ff7c82653e/cryptography-46.0.7-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:b7b412817be92117ec5ed95f880defe9cf18a832e8cafacf0a22337dc1981b4d", size = 4271354, upload-time = "2026-04-08T01:56:49.312Z" }, + { url = "https://files.pythonhosted.org/packages/50/46/cf71e26025c2e767c5609162c866a78e8a2915bbcfa408b7ca495c6140c4/cryptography-46.0.7-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:fbfd0e5f273877695cb93baf14b185f4878128b250cc9f8e617ea0c025dfb022", size = 4905845, upload-time = "2026-04-08T01:56:50.916Z" }, + { url = "https://files.pythonhosted.org/packages/c0/ea/01276740375bac6249d0a971ebdf6b4dc9ead0ee0a34ef3b5a88c1a9b0d4/cryptography-46.0.7-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:ffca7aa1d00cf7d6469b988c581598f2259e46215e0140af408966a24cf086ce", size = 4444641, upload-time = "2026-04-08T01:56:52.882Z" }, + { url = "https://files.pythonhosted.org/packages/3d/4c/7d258f169ae71230f25d9f3d06caabcff8c3baf0978e2b7d65e0acac3827/cryptography-46.0.7-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:60627cf07e0d9274338521205899337c5d18249db56865f943cbe753aa96f40f", size = 3967749, upload-time = "2026-04-08T01:56:54.597Z" }, + { url = "https://files.pythonhosted.org/packages/b5/2a/2ea0767cad19e71b3530e4cad9605d0b5e338b6a1e72c37c9c1ceb86c333/cryptography-46.0.7-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:80406c3065e2c55d7f49a9550fe0c49b3f12e5bfff5dedb727e319e1afb9bf99", size = 4270942, upload-time = "2026-04-08T01:56:56.416Z" }, + { url = "https://files.pythonhosted.org/packages/41/3d/fe14df95a83319af25717677e956567a105bb6ab25641acaa093db79975d/cryptography-46.0.7-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:c5b1ccd1239f48b7151a65bc6dd54bcfcc15e028c8ac126d3fada09db0e07ef1", size = 4871079, upload-time = "2026-04-08T01:56:58.31Z" }, + { url = "https://files.pythonhosted.org/packages/9c/59/4a479e0f36f8f378d397f4eab4c850b4ffb79a2f0d58704b8fa0703ddc11/cryptography-46.0.7-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:d5f7520159cd9c2154eb61eb67548ca05c5774d39e9c2c4339fd793fe7d097b2", size = 4443999, upload-time = "2026-04-08T01:57:00.508Z" }, + { url = "https://files.pythonhosted.org/packages/28/17/b59a741645822ec6d04732b43c5d35e4ef58be7bfa84a81e5ae6f05a1d33/cryptography-46.0.7-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:fcd8eac50d9138c1d7fc53a653ba60a2bee81a505f9f8850b6b2888555a45d0e", size = 4399191, upload-time = "2026-04-08T01:57:02.654Z" }, + { url = "https://files.pythonhosted.org/packages/59/6a/bb2e166d6d0e0955f1e9ff70f10ec4b2824c9cfcdb4da772c7dd69cc7d80/cryptography-46.0.7-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:65814c60f8cc400c63131584e3e1fad01235edba2614b61fbfbfa954082db0ee", size = 4655782, upload-time = "2026-04-08T01:57:04.592Z" }, + { url = "https://files.pythonhosted.org/packages/95/b6/3da51d48415bcb63b00dc17c2eff3a651b7c4fed484308d0f19b30e8cb2c/cryptography-46.0.7-cp314-cp314t-win32.whl", hash = "sha256:fdd1736fed309b4300346f88f74cd120c27c56852c3838cab416e7a166f67298", size = 3002227, upload-time = "2026-04-08T01:57:06.91Z" }, + { url = "https://files.pythonhosted.org/packages/32/a8/9f0e4ed57ec9cebe506e58db11ae472972ecb0c659e4d52bbaee80ca340a/cryptography-46.0.7-cp314-cp314t-win_amd64.whl", hash = "sha256:e06acf3c99be55aa3b516397fe42f5855597f430add9c17fa46bf2e0fb34c9bb", size = 3475332, upload-time = "2026-04-08T01:57:08.807Z" }, + { url = "https://files.pythonhosted.org/packages/a7/7f/cd42fc3614386bc0c12f0cb3c4ae1fc2bbca5c9662dfed031514911d513d/cryptography-46.0.7-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:462ad5cb1c148a22b2e3bcc5ad52504dff325d17daf5df8d88c17dda1f75f2a4", size = 7165618, upload-time = "2026-04-08T01:57:10.645Z" }, + { url = "https://files.pythonhosted.org/packages/a5/d0/36a49f0262d2319139d2829f773f1b97ef8aef7f97e6e5bd21455e5a8fb5/cryptography-46.0.7-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:84d4cced91f0f159a7ddacad249cc077e63195c36aac40b4150e7a57e84fffe7", size = 4270628, upload-time = "2026-04-08T01:57:12.885Z" }, + { url = "https://files.pythonhosted.org/packages/8a/6c/1a42450f464dda6ffbe578a911f773e54dd48c10f9895a23a7e88b3e7db5/cryptography-46.0.7-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:128c5edfe5e5938b86b03941e94fac9ee793a94452ad1365c9fc3f4f62216832", size = 4415405, upload-time = "2026-04-08T01:57:14.923Z" }, + { url = "https://files.pythonhosted.org/packages/9a/92/4ed714dbe93a066dc1f4b4581a464d2d7dbec9046f7c8b7016f5286329e2/cryptography-46.0.7-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:5e51be372b26ef4ba3de3c167cd3d1022934bc838ae9eaad7e644986d2a3d163", size = 4272715, upload-time = "2026-04-08T01:57:16.638Z" }, + { url = "https://files.pythonhosted.org/packages/b7/e6/a26b84096eddd51494bba19111f8fffe976f6a09f132706f8f1bf03f51f7/cryptography-46.0.7-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:cdf1a610ef82abb396451862739e3fc93b071c844399e15b90726ef7470eeaf2", size = 4918400, upload-time = "2026-04-08T01:57:19.021Z" }, + { url = "https://files.pythonhosted.org/packages/c7/08/ffd537b605568a148543ac3c2b239708ae0bd635064bab41359252ef88ed/cryptography-46.0.7-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1d25aee46d0c6f1a501adcddb2d2fee4b979381346a78558ed13e50aa8a59067", size = 4450634, upload-time = "2026-04-08T01:57:21.185Z" }, + { url = "https://files.pythonhosted.org/packages/16/01/0cd51dd86ab5b9befe0d031e276510491976c3a80e9f6e31810cce46c4ad/cryptography-46.0.7-cp38-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:cdfbe22376065ffcf8be74dc9a909f032df19bc58a699456a21712d6e5eabfd0", size = 3985233, upload-time = "2026-04-08T01:57:22.862Z" }, + { url = "https://files.pythonhosted.org/packages/92/49/819d6ed3a7d9349c2939f81b500a738cb733ab62fbecdbc1e38e83d45e12/cryptography-46.0.7-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:abad9dac36cbf55de6eb49badd4016806b3165d396f64925bf2999bcb67837ba", size = 4271955, upload-time = "2026-04-08T01:57:24.814Z" }, + { url = "https://files.pythonhosted.org/packages/80/07/ad9b3c56ebb95ed2473d46df0847357e01583f4c52a85754d1a55e29e4d0/cryptography-46.0.7-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:935ce7e3cfdb53e3536119a542b839bb94ec1ad081013e9ab9b7cfd478b05006", size = 4879888, upload-time = "2026-04-08T01:57:26.88Z" }, + { url = "https://files.pythonhosted.org/packages/b8/c7/201d3d58f30c4c2bdbe9b03844c291feb77c20511cc3586daf7edc12a47b/cryptography-46.0.7-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:35719dc79d4730d30f1c2b6474bd6acda36ae2dfae1e3c16f2051f215df33ce0", size = 4449961, upload-time = "2026-04-08T01:57:29.068Z" }, + { url = "https://files.pythonhosted.org/packages/a5/ef/649750cbf96f3033c3c976e112265c33906f8e462291a33d77f90356548c/cryptography-46.0.7-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7bbc6ccf49d05ac8f7d7b5e2e2c33830d4fe2061def88210a126d130d7f71a85", size = 4401696, upload-time = "2026-04-08T01:57:31.029Z" }, + { url = "https://files.pythonhosted.org/packages/41/52/a8908dcb1a389a459a29008c29966c1d552588d4ae6d43f3a1a4512e0ebe/cryptography-46.0.7-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a1529d614f44b863a7b480c6d000fe93b59acee9c82ffa027cfadc77521a9f5e", size = 4664256, upload-time = "2026-04-08T01:57:33.144Z" }, + { url = "https://files.pythonhosted.org/packages/4b/fa/f0ab06238e899cc3fb332623f337a7364f36f4bb3f2534c2bb95a35b132c/cryptography-46.0.7-cp38-abi3-win32.whl", hash = "sha256:f247c8c1a1fb45e12586afbb436ef21ff1e80670b2861a90353d9b025583d246", size = 3013001, upload-time = "2026-04-08T01:57:34.933Z" }, + { url = "https://files.pythonhosted.org/packages/d2/f1/00ce3bde3ca542d1acd8f8cfa38e446840945aa6363f9b74746394b14127/cryptography-46.0.7-cp38-abi3-win_amd64.whl", hash = "sha256:506c4ff91eff4f82bdac7633318a526b1d1309fc07ca76a3ad182cb5b686d6d3", size = 3472985, upload-time = "2026-04-08T01:57:36.714Z" }, + { url = "https://files.pythonhosted.org/packages/63/0c/dca8abb64e7ca4f6b2978769f6fea5ad06686a190cec381f0a796fdcaaba/cryptography-46.0.7-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:fc9ab8856ae6cf7c9358430e49b368f3108f050031442eaeb6b9d87e4dcf4e4f", size = 3476879, upload-time = "2026-04-08T01:57:38.664Z" }, + { url = "https://files.pythonhosted.org/packages/3a/ea/075aac6a84b7c271578d81a2f9968acb6e273002408729f2ddff517fed4a/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:d3b99c535a9de0adced13d159c5a9cf65c325601aa30f4be08afd680643e9c15", size = 4219700, upload-time = "2026-04-08T01:57:40.625Z" }, + { url = "https://files.pythonhosted.org/packages/6c/7b/1c55db7242b5e5612b29fc7a630e91ee7a6e3c8e7bf5406d22e206875fbd/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:d02c738dacda7dc2a74d1b2b3177042009d5cab7c7079db74afc19e56ca1b455", size = 4385982, upload-time = "2026-04-08T01:57:42.725Z" }, + { url = "https://files.pythonhosted.org/packages/cb/da/9870eec4b69c63ef5925bf7d8342b7e13bc2ee3d47791461c4e49ca212f4/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:04959522f938493042d595a736e7dbdff6eb6cc2339c11465b3ff89343b65f65", size = 4219115, upload-time = "2026-04-08T01:57:44.939Z" }, + { url = "https://files.pythonhosted.org/packages/f4/72/05aa5832b82dd341969e9a734d1812a6aadb088d9eb6f0430fc337cc5a8f/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:3986ac1dee6def53797289999eabe84798ad7817f3e97779b5061a95b0ee4968", size = 4385479, upload-time = "2026-04-08T01:57:46.86Z" }, + { url = "https://files.pythonhosted.org/packages/20/2a/1b016902351a523aa2bd446b50a5bc1175d7a7d1cf90fe2ef904f9b84ebc/cryptography-46.0.7-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:258514877e15963bd43b558917bc9f54cf7cf866c38aa576ebf47a77ddbc43a4", size = 3412829, upload-time = "2026-04-08T01:57:48.874Z" }, ] [[package]] @@ -339,7 +397,8 @@ source = { editable = "." } dependencies = [ { name = "cloudpickle" }, { name = "pyarrow" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] [package.dev-dependencies] @@ -492,11 +551,11 @@ wheels = [ [[package]] name = "idna" -version = "3.10" +version = "3.15" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f1/70/7703c29685631f5a7590aa73f1f1d3fa9a380e654b86af429e0934a32f7d/idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9", size = 190490, upload-time = "2024-09-15T18:07:39.745Z" } +sdist = { url = "https://files.pythonhosted.org/packages/82/77/7b3966d0b9d1d31a36ddf1746926a11dface89a83409bf1483f0237aa758/idna-3.15.tar.gz", hash = "sha256:ca962446ea538f7092a95e057da437618e886f4d349216d2b1e294abfdb65fdc", size = 199245, upload-time = "2026-05-12T22:45:57.011Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442, upload-time = "2024-09-15T18:07:37.964Z" }, + { url = "https://files.pythonhosted.org/packages/d2/23/408243171aa9aaba178d3e2559159c24c1171a641aa83b67bdd3394ead8e/idna-3.15-py3-none-any.whl", hash = "sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8", size = 72340, upload-time = "2026-05-12T22:45:55.733Z" }, ] [[package]] @@ -532,7 +591,8 @@ dependencies = [ { name = "pygments" }, { name = "stack-data" }, { name = "traitlets" }, - { name = "typing-extensions", marker = "python_full_version < '3.12'" }, + { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/01/35/6f90fdddff7a08b7b715fccbd2427b5212c9525cd043d26fdc45bee0708d/ipython-8.31.0.tar.gz", hash = "sha256:b6a2274606bec6166405ff05e54932ed6e5cfecaca1fc05f2cacde7bb074d70b", size = 5501011, upload-time = "2024-12-20T12:34:22.61Z" } wheels = [ @@ -1091,59 +1151,59 @@ wheels = [ [[package]] name = "pyarrow" -version = "22.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/30/53/04a7fdc63e6056116c9ddc8b43bc28c12cdd181b85cbeadb79278475f3ae/pyarrow-22.0.0.tar.gz", hash = "sha256:3d600dc583260d845c7d8a6db540339dd883081925da2bd1c5cb808f720b3cd9", size = 1151151, upload-time = "2025-10-24T12:30:00.762Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d9/9b/cb3f7e0a345353def531ca879053e9ef6b9f38ed91aebcf68b09ba54dec0/pyarrow-22.0.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:77718810bd3066158db1e95a63c160ad7ce08c6b0710bc656055033e39cdad88", size = 34223968, upload-time = "2025-10-24T10:03:31.21Z" }, - { url = "https://files.pythonhosted.org/packages/6c/41/3184b8192a120306270c5307f105b70320fdaa592c99843c5ef78aaefdcf/pyarrow-22.0.0-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:44d2d26cda26d18f7af7db71453b7b783788322d756e81730acb98f24eb90ace", size = 35942085, upload-time = "2025-10-24T10:03:38.146Z" }, - { url = "https://files.pythonhosted.org/packages/d9/3d/a1eab2f6f08001f9fb714b8ed5cfb045e2fe3e3e3c0c221f2c9ed1e6d67d/pyarrow-22.0.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:b9d71701ce97c95480fecb0039ec5bb889e75f110da72005743451339262f4ce", size = 44964613, upload-time = "2025-10-24T10:03:46.516Z" }, - { url = "https://files.pythonhosted.org/packages/46/46/a1d9c24baf21cfd9ce994ac820a24608decf2710521b29223d4334985127/pyarrow-22.0.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:710624ab925dc2b05a6229d47f6f0dac1c1155e6ed559be7109f684eba048a48", size = 47627059, upload-time = "2025-10-24T10:03:55.353Z" }, - { url = "https://files.pythonhosted.org/packages/3a/4c/f711acb13075c1391fd54bc17e078587672c575f8de2a6e62509af026dcf/pyarrow-22.0.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:f963ba8c3b0199f9d6b794c90ec77545e05eadc83973897a4523c9e8d84e9340", size = 47947043, upload-time = "2025-10-24T10:04:05.408Z" }, - { url = "https://files.pythonhosted.org/packages/4e/70/1f3180dd7c2eab35c2aca2b29ace6c519f827dcd4cfeb8e0dca41612cf7a/pyarrow-22.0.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:bd0d42297ace400d8febe55f13fdf46e86754842b860c978dfec16f081e5c653", size = 50206505, upload-time = "2025-10-24T10:04:15.786Z" }, - { url = "https://files.pythonhosted.org/packages/80/07/fea6578112c8c60ffde55883a571e4c4c6bc7049f119d6b09333b5cc6f73/pyarrow-22.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:00626d9dc0f5ef3a75fe63fd68b9c7c8302d2b5bbc7f74ecaedba83447a24f84", size = 28101641, upload-time = "2025-10-24T10:04:22.57Z" }, - { url = "https://files.pythonhosted.org/packages/2e/b7/18f611a8cdc43417f9394a3ccd3eace2f32183c08b9eddc3d17681819f37/pyarrow-22.0.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:3e294c5eadfb93d78b0763e859a0c16d4051fc1c5231ae8956d61cb0b5666f5a", size = 34272022, upload-time = "2025-10-24T10:04:28.973Z" }, - { url = "https://files.pythonhosted.org/packages/26/5c/f259e2526c67eb4b9e511741b19870a02363a47a35edbebc55c3178db22d/pyarrow-22.0.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:69763ab2445f632d90b504a815a2a033f74332997052b721002298ed6de40f2e", size = 35995834, upload-time = "2025-10-24T10:04:35.467Z" }, - { url = "https://files.pythonhosted.org/packages/50/8d/281f0f9b9376d4b7f146913b26fac0aa2829cd1ee7e997f53a27411bbb92/pyarrow-22.0.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:b41f37cabfe2463232684de44bad753d6be08a7a072f6a83447eeaf0e4d2a215", size = 45030348, upload-time = "2025-10-24T10:04:43.366Z" }, - { url = "https://files.pythonhosted.org/packages/f5/e5/53c0a1c428f0976bf22f513d79c73000926cb00b9c138d8e02daf2102e18/pyarrow-22.0.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:35ad0f0378c9359b3f297299c3309778bb03b8612f987399a0333a560b43862d", size = 47699480, upload-time = "2025-10-24T10:04:51.486Z" }, - { url = "https://files.pythonhosted.org/packages/95/e1/9dbe4c465c3365959d183e6345d0a8d1dc5b02ca3f8db4760b3bc834cf25/pyarrow-22.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8382ad21458075c2e66a82a29d650f963ce51c7708c7c0ff313a8c206c4fd5e8", size = 48011148, upload-time = "2025-10-24T10:04:59.585Z" }, - { url = "https://files.pythonhosted.org/packages/c5/b4/7caf5d21930061444c3cf4fa7535c82faf5263e22ce43af7c2759ceb5b8b/pyarrow-22.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1a812a5b727bc09c3d7ea072c4eebf657c2f7066155506ba31ebf4792f88f016", size = 50276964, upload-time = "2025-10-24T10:05:08.175Z" }, - { url = "https://files.pythonhosted.org/packages/ae/f3/cec89bd99fa3abf826f14d4e53d3d11340ce6f6af4d14bdcd54cd83b6576/pyarrow-22.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:ec5d40dd494882704fb876c16fa7261a69791e784ae34e6b5992e977bd2e238c", size = 28106517, upload-time = "2025-10-24T10:05:14.314Z" }, - { url = "https://files.pythonhosted.org/packages/af/63/ba23862d69652f85b615ca14ad14f3bcfc5bf1b99ef3f0cd04ff93fdad5a/pyarrow-22.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:bea79263d55c24a32b0d79c00a1c58bb2ee5f0757ed95656b01c0fb310c5af3d", size = 34211578, upload-time = "2025-10-24T10:05:21.583Z" }, - { url = "https://files.pythonhosted.org/packages/b1/d0/f9ad86fe809efd2bcc8be32032fa72e8b0d112b01ae56a053006376c5930/pyarrow-22.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:12fe549c9b10ac98c91cf791d2945e878875d95508e1a5d14091a7aaa66d9cf8", size = 35989906, upload-time = "2025-10-24T10:05:29.485Z" }, - { url = "https://files.pythonhosted.org/packages/b4/a8/f910afcb14630e64d673f15904ec27dd31f1e009b77033c365c84e8c1e1d/pyarrow-22.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:334f900ff08ce0423407af97e6c26ad5d4e3b0763645559ece6fbf3747d6a8f5", size = 45021677, upload-time = "2025-10-24T10:05:38.274Z" }, - { url = "https://files.pythonhosted.org/packages/13/95/aec81f781c75cd10554dc17a25849c720d54feafb6f7847690478dcf5ef8/pyarrow-22.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:c6c791b09c57ed76a18b03f2631753a4960eefbbca80f846da8baefc6491fcfe", size = 47726315, upload-time = "2025-10-24T10:05:47.314Z" }, - { url = "https://files.pythonhosted.org/packages/bb/d4/74ac9f7a54cfde12ee42734ea25d5a3c9a45db78f9def949307a92720d37/pyarrow-22.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c3200cb41cdbc65156e5f8c908d739b0dfed57e890329413da2748d1a2cd1a4e", size = 47990906, upload-time = "2025-10-24T10:05:58.254Z" }, - { url = "https://files.pythonhosted.org/packages/2e/71/fedf2499bf7a95062eafc989ace56572f3343432570e1c54e6599d5b88da/pyarrow-22.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ac93252226cf288753d8b46280f4edf3433bf9508b6977f8dd8526b521a1bbb9", size = 50306783, upload-time = "2025-10-24T10:06:08.08Z" }, - { url = "https://files.pythonhosted.org/packages/68/ed/b202abd5a5b78f519722f3d29063dda03c114711093c1995a33b8e2e0f4b/pyarrow-22.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:44729980b6c50a5f2bfcc2668d36c569ce17f8b17bccaf470c4313dcbbf13c9d", size = 27972883, upload-time = "2025-10-24T10:06:14.204Z" }, - { url = "https://files.pythonhosted.org/packages/a6/d6/d0fac16a2963002fc22c8fa75180a838737203d558f0ed3b564c4a54eef5/pyarrow-22.0.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:e6e95176209257803a8b3d0394f21604e796dadb643d2f7ca21b66c9c0b30c9a", size = 34204629, upload-time = "2025-10-24T10:06:20.274Z" }, - { url = "https://files.pythonhosted.org/packages/c6/9c/1d6357347fbae062ad3f17082f9ebc29cc733321e892c0d2085f42a2212b/pyarrow-22.0.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:001ea83a58024818826a9e3f89bf9310a114f7e26dfe404a4c32686f97bd7901", size = 35985783, upload-time = "2025-10-24T10:06:27.301Z" }, - { url = "https://files.pythonhosted.org/packages/ff/c0/782344c2ce58afbea010150df07e3a2f5fdad299cd631697ae7bd3bac6e3/pyarrow-22.0.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:ce20fe000754f477c8a9125543f1936ea5b8867c5406757c224d745ed033e691", size = 45020999, upload-time = "2025-10-24T10:06:35.387Z" }, - { url = "https://files.pythonhosted.org/packages/1b/8b/5362443737a5307a7b67c1017c42cd104213189b4970bf607e05faf9c525/pyarrow-22.0.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:e0a15757fccb38c410947df156f9749ae4a3c89b2393741a50521f39a8cf202a", size = 47724601, upload-time = "2025-10-24T10:06:43.551Z" }, - { url = "https://files.pythonhosted.org/packages/69/4d/76e567a4fc2e190ee6072967cb4672b7d9249ac59ae65af2d7e3047afa3b/pyarrow-22.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cedb9dd9358e4ea1d9bce3665ce0797f6adf97ff142c8e25b46ba9cdd508e9b6", size = 48001050, upload-time = "2025-10-24T10:06:52.284Z" }, - { url = "https://files.pythonhosted.org/packages/01/5e/5653f0535d2a1aef8223cee9d92944cb6bccfee5cf1cd3f462d7cb022790/pyarrow-22.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:252be4a05f9d9185bb8c18e83764ebcfea7185076c07a7a662253af3a8c07941", size = 50307877, upload-time = "2025-10-24T10:07:02.405Z" }, - { url = "https://files.pythonhosted.org/packages/2d/f8/1d0bd75bf9328a3b826e24a16e5517cd7f9fbf8d34a3184a4566ef5a7f29/pyarrow-22.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:a4893d31e5ef780b6edcaf63122df0f8d321088bb0dee4c8c06eccb1ca28d145", size = 27977099, upload-time = "2025-10-24T10:08:07.259Z" }, - { url = "https://files.pythonhosted.org/packages/90/81/db56870c997805bf2b0f6eeeb2d68458bf4654652dccdcf1bf7a42d80903/pyarrow-22.0.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:f7fe3dbe871294ba70d789be16b6e7e52b418311e166e0e3cba9522f0f437fb1", size = 34336685, upload-time = "2025-10-24T10:07:11.47Z" }, - { url = "https://files.pythonhosted.org/packages/1c/98/0727947f199aba8a120f47dfc229eeb05df15bcd7a6f1b669e9f882afc58/pyarrow-22.0.0-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:ba95112d15fd4f1105fb2402c4eab9068f0554435e9b7085924bcfaac2cc306f", size = 36032158, upload-time = "2025-10-24T10:07:18.626Z" }, - { url = "https://files.pythonhosted.org/packages/96/b4/9babdef9c01720a0785945c7cf550e4acd0ebcd7bdd2e6f0aa7981fa85e2/pyarrow-22.0.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:c064e28361c05d72eed8e744c9605cbd6d2bb7481a511c74071fd9b24bc65d7d", size = 44892060, upload-time = "2025-10-24T10:07:26.002Z" }, - { url = "https://files.pythonhosted.org/packages/f8/ca/2f8804edd6279f78a37062d813de3f16f29183874447ef6d1aadbb4efa0f/pyarrow-22.0.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6f9762274496c244d951c819348afbcf212714902742225f649cf02823a6a10f", size = 47504395, upload-time = "2025-10-24T10:07:34.09Z" }, - { url = "https://files.pythonhosted.org/packages/b9/f0/77aa5198fd3943682b2e4faaf179a674f0edea0d55d326d83cb2277d9363/pyarrow-22.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:a9d9ffdc2ab696f6b15b4d1f7cec6658e1d788124418cb30030afbae31c64746", size = 48066216, upload-time = "2025-10-24T10:07:43.528Z" }, - { url = "https://files.pythonhosted.org/packages/79/87/a1937b6e78b2aff18b706d738c9e46ade5bfcf11b294e39c87706a0089ac/pyarrow-22.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:ec1a15968a9d80da01e1d30349b2b0d7cc91e96588ee324ce1b5228175043e95", size = 50288552, upload-time = "2025-10-24T10:07:53.519Z" }, - { url = "https://files.pythonhosted.org/packages/60/ae/b5a5811e11f25788ccfdaa8f26b6791c9807119dffcf80514505527c384c/pyarrow-22.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:bba208d9c7decf9961998edf5c65e3ea4355d5818dd6cd0f6809bec1afb951cc", size = 28262504, upload-time = "2025-10-24T10:08:00.932Z" }, - { url = "https://files.pythonhosted.org/packages/bd/b0/0fa4d28a8edb42b0a7144edd20befd04173ac79819547216f8a9f36f9e50/pyarrow-22.0.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:9bddc2cade6561f6820d4cd73f99a0243532ad506bc510a75a5a65a522b2d74d", size = 34224062, upload-time = "2025-10-24T10:08:14.101Z" }, - { url = "https://files.pythonhosted.org/packages/0f/a8/7a719076b3c1be0acef56a07220c586f25cd24de0e3f3102b438d18ae5df/pyarrow-22.0.0-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:e70ff90c64419709d38c8932ea9fe1cc98415c4f87ea8da81719e43f02534bc9", size = 35990057, upload-time = "2025-10-24T10:08:21.842Z" }, - { url = "https://files.pythonhosted.org/packages/89/3c/359ed54c93b47fb6fe30ed16cdf50e3f0e8b9ccfb11b86218c3619ae50a8/pyarrow-22.0.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:92843c305330aa94a36e706c16209cd4df274693e777ca47112617db7d0ef3d7", size = 45068002, upload-time = "2025-10-24T10:08:29.034Z" }, - { url = "https://files.pythonhosted.org/packages/55/fc/4945896cc8638536ee787a3bd6ce7cec8ec9acf452d78ec39ab328efa0a1/pyarrow-22.0.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:6dda1ddac033d27421c20d7a7943eec60be44e0db4e079f33cc5af3b8280ccde", size = 47737765, upload-time = "2025-10-24T10:08:38.559Z" }, - { url = "https://files.pythonhosted.org/packages/cd/5e/7cb7edeb2abfaa1f79b5d5eb89432356155c8426f75d3753cbcb9592c0fd/pyarrow-22.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:84378110dd9a6c06323b41b56e129c504d157d1a983ce8f5443761eb5256bafc", size = 48048139, upload-time = "2025-10-24T10:08:46.784Z" }, - { url = "https://files.pythonhosted.org/packages/88/c6/546baa7c48185f5e9d6e59277c4b19f30f48c94d9dd938c2a80d4d6b067c/pyarrow-22.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:854794239111d2b88b40b6ef92aa478024d1e5074f364033e73e21e3f76b25e0", size = 50314244, upload-time = "2025-10-24T10:08:55.771Z" }, - { url = "https://files.pythonhosted.org/packages/3c/79/755ff2d145aafec8d347bf18f95e4e81c00127f06d080135dfc86aea417c/pyarrow-22.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:b883fe6fd85adad7932b3271c38ac289c65b7337c2c132e9569f9d3940620730", size = 28757501, upload-time = "2025-10-24T10:09:59.891Z" }, - { url = "https://files.pythonhosted.org/packages/0e/d2/237d75ac28ced3147912954e3c1a174df43a95f4f88e467809118a8165e0/pyarrow-22.0.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:7a820d8ae11facf32585507c11f04e3f38343c1e784c9b5a8b1da5c930547fe2", size = 34355506, upload-time = "2025-10-24T10:09:02.953Z" }, - { url = "https://files.pythonhosted.org/packages/1e/2c/733dfffe6d3069740f98e57ff81007809067d68626c5faef293434d11bd6/pyarrow-22.0.0-cp314-cp314t-macosx_12_0_x86_64.whl", hash = "sha256:c6ec3675d98915bf1ec8b3c7986422682f7232ea76cad276f4c8abd5b7319b70", size = 36047312, upload-time = "2025-10-24T10:09:10.334Z" }, - { url = "https://files.pythonhosted.org/packages/7c/2b/29d6e3782dc1f299727462c1543af357a0f2c1d3c160ce199950d9ca51eb/pyarrow-22.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:3e739edd001b04f654b166204fc7a9de896cf6007eaff33409ee9e50ceaff754", size = 45081609, upload-time = "2025-10-24T10:09:18.61Z" }, - { url = "https://files.pythonhosted.org/packages/8d/42/aa9355ecc05997915af1b7b947a7f66c02dcaa927f3203b87871c114ba10/pyarrow-22.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:7388ac685cab5b279a41dfe0a6ccd99e4dbf322edfb63e02fc0443bf24134e91", size = 47703663, upload-time = "2025-10-24T10:09:27.369Z" }, - { url = "https://files.pythonhosted.org/packages/ee/62/45abedde480168e83a1de005b7b7043fd553321c1e8c5a9a114425f64842/pyarrow-22.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f633074f36dbc33d5c05b5dc75371e5660f1dbf9c8b1d95669def05e5425989c", size = 48066543, upload-time = "2025-10-24T10:09:34.908Z" }, - { url = "https://files.pythonhosted.org/packages/84/e9/7878940a5b072e4f3bf998770acafeae13b267f9893af5f6d4ab3904b67e/pyarrow-22.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:4c19236ae2402a8663a2c8f21f1870a03cc57f0bef7e4b6eb3238cc82944de80", size = 50288838, upload-time = "2025-10-24T10:09:44.394Z" }, - { url = "https://files.pythonhosted.org/packages/7b/03/f335d6c52b4a4761bcc83499789a1e2e16d9d201a58c327a9b5cc9a41bd9/pyarrow-22.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:0c34fe18094686194f204a3b1787a27456897d8a2d62caf84b61e8dfbc0252ae", size = 29185594, upload-time = "2025-10-24T10:09:53.111Z" }, +version = "23.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/88/22/134986a4cc224d593c1afde5494d18ff629393d74cc2eddb176669f234a4/pyarrow-23.0.1.tar.gz", hash = "sha256:b8c5873e33440b2bc2f4a79d2b47017a89c5a24116c055625e6f2ee50523f019", size = 1167336, upload-time = "2026-02-16T10:14:12.39Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/a8/24e5dc6855f50a62936ceb004e6e9645e4219a8065f304145d7fb8a79d5d/pyarrow-23.0.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:3fab8f82571844eb3c460f90a75583801d14ca0cc32b1acc8c361650e006fd56", size = 34307390, upload-time = "2026-02-16T10:08:08.654Z" }, + { url = "https://files.pythonhosted.org/packages/bc/8e/4be5617b4aaae0287f621ad31c6036e5f63118cfca0dc57d42121ff49b51/pyarrow-23.0.1-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:3f91c038b95f71ddfc865f11d5876c42f343b4495535bd262c7b321b0b94507c", size = 35853761, upload-time = "2026-02-16T10:08:17.811Z" }, + { url = "https://files.pythonhosted.org/packages/2e/08/3e56a18819462210432ae37d10f5c8eed3828be1d6c751b6e6a2e93c286a/pyarrow-23.0.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:d0744403adabef53c985a7f8a082b502a368510c40d184df349a0a8754533258", size = 44493116, upload-time = "2026-02-16T10:08:25.792Z" }, + { url = "https://files.pythonhosted.org/packages/f8/82/c40b68001dbec8a3faa4c08cd8c200798ac732d2854537c5449dc859f55a/pyarrow-23.0.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:c33b5bf406284fd0bba436ed6f6c3ebe8e311722b441d89397c54f871c6863a2", size = 47564532, upload-time = "2026-02-16T10:08:34.27Z" }, + { url = "https://files.pythonhosted.org/packages/20/bc/73f611989116b6f53347581b02177f9f620efdf3cd3f405d0e83cdf53a83/pyarrow-23.0.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:ddf743e82f69dcd6dbbcb63628895d7161e04e56794ef80550ac6f3315eeb1d5", size = 48183685, upload-time = "2026-02-16T10:08:42.889Z" }, + { url = "https://files.pythonhosted.org/packages/b0/cc/6c6b3ecdae2a8c3aced99956187e8302fc954cc2cca2a37cf2111dad16ce/pyarrow-23.0.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:e052a211c5ac9848ae15d5ec875ed0943c0221e2fcfe69eee80b604b4e703222", size = 50605582, upload-time = "2026-02-16T10:08:51.641Z" }, + { url = "https://files.pythonhosted.org/packages/8d/94/d359e708672878d7638a04a0448edf7c707f9e5606cee11e15aaa5c7535a/pyarrow-23.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:5abde149bb3ce524782d838eb67ac095cd3fd6090eba051130589793f1a7f76d", size = 27521148, upload-time = "2026-02-16T10:08:58.077Z" }, + { url = "https://files.pythonhosted.org/packages/b0/41/8e6b6ef7e225d4ceead8459427a52afdc23379768f54dd3566014d7618c1/pyarrow-23.0.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:6f0147ee9e0386f519c952cc670eb4a8b05caa594eeffe01af0e25f699e4e9bb", size = 34302230, upload-time = "2026-02-16T10:09:03.859Z" }, + { url = "https://files.pythonhosted.org/packages/bf/4a/1472c00392f521fea03ae93408bf445cc7bfa1ab81683faf9bc188e36629/pyarrow-23.0.1-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:0ae6e17c828455b6265d590100c295193f93cc5675eb0af59e49dbd00d2de350", size = 35850050, upload-time = "2026-02-16T10:09:11.877Z" }, + { url = "https://files.pythonhosted.org/packages/0c/b2/bd1f2f05ded56af7f54d702c8364c9c43cd6abb91b0e9933f3d77b4f4132/pyarrow-23.0.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:fed7020203e9ef273360b9e45be52a2a47d3103caf156a30ace5247ffb51bdbd", size = 44491918, upload-time = "2026-02-16T10:09:18.144Z" }, + { url = "https://files.pythonhosted.org/packages/0b/62/96459ef5b67957eac38a90f541d1c28833d1b367f014a482cb63f3b7cd2d/pyarrow-23.0.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:26d50dee49d741ac0e82185033488d28d35be4d763ae6f321f97d1140eb7a0e9", size = 47562811, upload-time = "2026-02-16T10:09:25.792Z" }, + { url = "https://files.pythonhosted.org/packages/7d/94/1170e235add1f5f45a954e26cd0e906e7e74e23392dcb560de471f7366ec/pyarrow-23.0.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:3c30143b17161310f151f4a2bcfe41b5ff744238c1039338779424e38579d701", size = 48183766, upload-time = "2026-02-16T10:09:34.645Z" }, + { url = "https://files.pythonhosted.org/packages/0e/2d/39a42af4570377b99774cdb47f63ee6c7da7616bd55b3d5001aa18edfe4f/pyarrow-23.0.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:db2190fa79c80a23fdd29fef4b8992893f024ae7c17d2f5f4db7171fa30c2c78", size = 50607669, upload-time = "2026-02-16T10:09:44.153Z" }, + { url = "https://files.pythonhosted.org/packages/00/ca/db94101c187f3df742133ac837e93b1f269ebdac49427f8310ee40b6a58f/pyarrow-23.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:f00f993a8179e0e1c9713bcc0baf6d6c01326a406a9c23495ec1ba9c9ebf2919", size = 27527698, upload-time = "2026-02-16T10:09:50.263Z" }, + { url = "https://files.pythonhosted.org/packages/9a/4b/4166bb5abbfe6f750fc60ad337c43ecf61340fa52ab386da6e8dbf9e63c4/pyarrow-23.0.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:f4b0dbfa124c0bb161f8b5ebb40f1a680b70279aa0c9901d44a2b5a20806039f", size = 34214575, upload-time = "2026-02-16T10:09:56.225Z" }, + { url = "https://files.pythonhosted.org/packages/e1/da/3f941e3734ac8088ea588b53e860baeddac8323ea40ce22e3d0baa865cc9/pyarrow-23.0.1-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:7707d2b6673f7de054e2e83d59f9e805939038eebe1763fe811ee8fa5c0cd1a7", size = 35832540, upload-time = "2026-02-16T10:10:03.428Z" }, + { url = "https://files.pythonhosted.org/packages/88/7c/3d841c366620e906d54430817531b877ba646310296df42ef697308c2705/pyarrow-23.0.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:86ff03fb9f1a320266e0de855dee4b17da6794c595d207f89bba40d16b5c78b9", size = 44470940, upload-time = "2026-02-16T10:10:10.704Z" }, + { url = "https://files.pythonhosted.org/packages/2c/a5/da83046273d990f256cb79796a190bbf7ec999269705ddc609403f8c6b06/pyarrow-23.0.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:813d99f31275919c383aab17f0f455a04f5a429c261cc411b1e9a8f5e4aaaa05", size = 47586063, upload-time = "2026-02-16T10:10:17.95Z" }, + { url = "https://files.pythonhosted.org/packages/5b/3c/b7d2ebcff47a514f47f9da1e74b7949138c58cfeb108cdd4ee62f43f0cf3/pyarrow-23.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bf5842f960cddd2ef757d486041d57c96483efc295a8c4a0e20e704cbbf39c67", size = 48173045, upload-time = "2026-02-16T10:10:25.363Z" }, + { url = "https://files.pythonhosted.org/packages/43/b2/b40961262213beaba6acfc88698eb773dfce32ecdf34d19291db94c2bd73/pyarrow-23.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:564baf97c858ecc03ec01a41062e8f4698abc3e6e2acd79c01c2e97880a19730", size = 50621741, upload-time = "2026-02-16T10:10:33.477Z" }, + { url = "https://files.pythonhosted.org/packages/f6/70/1fdda42d65b28b078e93d75d371b2185a61da89dda4def8ba6ba41ebdeb4/pyarrow-23.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:07deae7783782ac7250989a7b2ecde9b3c343a643f82e8a4df03d93b633006f0", size = 27620678, upload-time = "2026-02-16T10:10:39.31Z" }, + { url = "https://files.pythonhosted.org/packages/47/10/2cbe4c6f0fb83d2de37249567373d64327a5e4d8db72f486db42875b08f6/pyarrow-23.0.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:6b8fda694640b00e8af3c824f99f789e836720aa8c9379fb435d4c4953a756b8", size = 34210066, upload-time = "2026-02-16T10:10:45.487Z" }, + { url = "https://files.pythonhosted.org/packages/cb/4f/679fa7e84dadbaca7a65f7cdba8d6c83febbd93ca12fa4adf40ba3b6362b/pyarrow-23.0.1-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:8ff51b1addc469b9444b7c6f3548e19dc931b172ab234e995a60aea9f6e6025f", size = 35825526, upload-time = "2026-02-16T10:10:52.266Z" }, + { url = "https://files.pythonhosted.org/packages/f9/63/d2747d930882c9d661e9398eefc54f15696547b8983aaaf11d4a2e8b5426/pyarrow-23.0.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:71c5be5cbf1e1cb6169d2a0980850bccb558ddc9b747b6206435313c47c37677", size = 44473279, upload-time = "2026-02-16T10:11:01.557Z" }, + { url = "https://files.pythonhosted.org/packages/b3/93/10a48b5e238de6d562a411af6467e71e7aedbc9b87f8d3a35f1560ae30fb/pyarrow-23.0.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:9b6f4f17b43bc39d56fec96e53fe89d94bac3eb134137964371b45352d40d0c2", size = 47585798, upload-time = "2026-02-16T10:11:09.401Z" }, + { url = "https://files.pythonhosted.org/packages/5c/20/476943001c54ef078dbf9542280e22741219a184a0632862bca4feccd666/pyarrow-23.0.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:9fc13fc6c403d1337acab46a2c4346ca6c9dec5780c3c697cf8abfd5e19b6b37", size = 48179446, upload-time = "2026-02-16T10:11:17.781Z" }, + { url = "https://files.pythonhosted.org/packages/4b/b6/5dd0c47b335fcd8edba9bfab78ad961bd0fd55ebe53468cc393f45e0be60/pyarrow-23.0.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5c16ed4f53247fa3ffb12a14d236de4213a4415d127fe9cebed33d51671113e2", size = 50623972, upload-time = "2026-02-16T10:11:26.185Z" }, + { url = "https://files.pythonhosted.org/packages/d5/09/a532297c9591a727d67760e2e756b83905dd89adb365a7f6e9c72578bcc1/pyarrow-23.0.1-cp313-cp313-win_amd64.whl", hash = "sha256:cecfb12ef629cf6be0b1887f9f86463b0dd3dc3195ae6224e74006be4736035a", size = 27540749, upload-time = "2026-02-16T10:12:23.297Z" }, + { url = "https://files.pythonhosted.org/packages/a5/8e/38749c4b1303e6ae76b3c80618f84861ae0c55dd3c2273842ea6f8258233/pyarrow-23.0.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:29f7f7419a0e30264ea261fdc0e5fe63ce5a6095003db2945d7cd78df391a7e1", size = 34471544, upload-time = "2026-02-16T10:11:32.535Z" }, + { url = "https://files.pythonhosted.org/packages/a3/73/f237b2bc8c669212f842bcfd842b04fc8d936bfc9d471630569132dc920d/pyarrow-23.0.1-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:33d648dc25b51fd8055c19e4261e813dfc4d2427f068bcecc8b53d01b81b0500", size = 35949911, upload-time = "2026-02-16T10:11:39.813Z" }, + { url = "https://files.pythonhosted.org/packages/0c/86/b912195eee0903b5611bf596833def7d146ab2d301afeb4b722c57ffc966/pyarrow-23.0.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:cd395abf8f91c673dd3589cadc8cc1ee4e8674fa61b2e923c8dd215d9c7d1f41", size = 44520337, upload-time = "2026-02-16T10:11:47.764Z" }, + { url = "https://files.pythonhosted.org/packages/69/c2/f2a717fb824f62d0be952ea724b4f6f9372a17eed6f704b5c9526f12f2f1/pyarrow-23.0.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:00be9576d970c31defb5c32eb72ef585bf600ef6d0a82d5eccaae96639cf9d07", size = 47548944, upload-time = "2026-02-16T10:11:56.607Z" }, + { url = "https://files.pythonhosted.org/packages/84/a7/90007d476b9f0dc308e3bc57b832d004f848fd6c0da601375d20d92d1519/pyarrow-23.0.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c2139549494445609f35a5cda4eb94e2c9e4d704ce60a095b342f82460c73a83", size = 48236269, upload-time = "2026-02-16T10:12:04.47Z" }, + { url = "https://files.pythonhosted.org/packages/b0/3f/b16fab3e77709856eb6ac328ce35f57a6d4a18462c7ca5186ef31b45e0e0/pyarrow-23.0.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:7044b442f184d84e2351e5084600f0d7343d6117aabcbc1ac78eb1ae11eb4125", size = 50604794, upload-time = "2026-02-16T10:12:11.797Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a1/22df0620a9fac31d68397a75465c344e83c3dfe521f7612aea33e27ab6c0/pyarrow-23.0.1-cp313-cp313t-win_amd64.whl", hash = "sha256:a35581e856a2fafa12f3f54fce4331862b1cfb0bef5758347a858a4aa9d6bae8", size = 27660642, upload-time = "2026-02-16T10:12:17.746Z" }, + { url = "https://files.pythonhosted.org/packages/8d/1b/6da9a89583ce7b23ac611f183ae4843cd3a6cf54f079549b0e8c14031e73/pyarrow-23.0.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:5df1161da23636a70838099d4aaa65142777185cc0cdba4037a18cee7d8db9ca", size = 34238755, upload-time = "2026-02-16T10:12:32.819Z" }, + { url = "https://files.pythonhosted.org/packages/ae/b5/d58a241fbe324dbaeb8df07be6af8752c846192d78d2272e551098f74e88/pyarrow-23.0.1-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:fa8e51cb04b9f8c9c5ace6bab63af9a1f88d35c0d6cbf53e8c17c098552285e1", size = 35847826, upload-time = "2026-02-16T10:12:38.949Z" }, + { url = "https://files.pythonhosted.org/packages/54/a5/8cbc83f04aba433ca7b331b38f39e000efd9f0c7ce47128670e737542996/pyarrow-23.0.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:0b95a3994f015be13c63148fef8832e8a23938128c185ee951c98908a696e0eb", size = 44536859, upload-time = "2026-02-16T10:12:45.467Z" }, + { url = "https://files.pythonhosted.org/packages/36/2e/c0f017c405fcdc252dbccafbe05e36b0d0eb1ea9a958f081e01c6972927f/pyarrow-23.0.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:4982d71350b1a6e5cfe1af742c53dfb759b11ce14141870d05d9e540d13bc5d1", size = 47614443, upload-time = "2026-02-16T10:12:55.525Z" }, + { url = "https://files.pythonhosted.org/packages/af/6b/2314a78057912f5627afa13ba43809d9d653e6630859618b0fd81a4e0759/pyarrow-23.0.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c250248f1fe266db627921c89b47b7c06fee0489ad95b04d50353537d74d6886", size = 48232991, upload-time = "2026-02-16T10:13:04.729Z" }, + { url = "https://files.pythonhosted.org/packages/40/f2/1bcb1d3be3460832ef3370d621142216e15a2c7c62602a4ea19ec240dd64/pyarrow-23.0.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5f4763b83c11c16e5f4c15601ba6dfa849e20723b46aa2617cb4bffe8768479f", size = 50645077, upload-time = "2026-02-16T10:13:14.147Z" }, + { url = "https://files.pythonhosted.org/packages/eb/3f/b1da7b61cd66566a4d4c8383d376c606d1c34a906c3f1cb35c479f59d1aa/pyarrow-23.0.1-cp314-cp314-win_amd64.whl", hash = "sha256:3a4c85ef66c134161987c17b147d6bffdca4566f9a4c1d81a0a01cdf08414ea5", size = 28234271, upload-time = "2026-02-16T10:14:09.397Z" }, + { url = "https://files.pythonhosted.org/packages/b5/78/07f67434e910a0f7323269be7bfbf58699bd0c1d080b18a1ab49ba943fe8/pyarrow-23.0.1-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:17cd28e906c18af486a499422740298c52d7c6795344ea5002a7720b4eadf16d", size = 34488692, upload-time = "2026-02-16T10:13:21.541Z" }, + { url = "https://files.pythonhosted.org/packages/50/76/34cf7ae93ece1f740a04910d9f7e80ba166b9b4ab9596a953e9e62b90fe1/pyarrow-23.0.1-cp314-cp314t-macosx_12_0_x86_64.whl", hash = "sha256:76e823d0e86b4fb5e1cf4a58d293036e678b5a4b03539be933d3b31f9406859f", size = 35964383, upload-time = "2026-02-16T10:13:28.63Z" }, + { url = "https://files.pythonhosted.org/packages/46/90/459b827238936d4244214be7c684e1b366a63f8c78c380807ae25ed92199/pyarrow-23.0.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:a62e1899e3078bf65943078b3ad2a6ddcacf2373bc06379aac61b1e548a75814", size = 44538119, upload-time = "2026-02-16T10:13:35.506Z" }, + { url = "https://files.pythonhosted.org/packages/28/a1/93a71ae5881e99d1f9de1d4554a87be37da11cd6b152239fb5bd924fdc64/pyarrow-23.0.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:df088e8f640c9fae3b1f495b3c64755c4e719091caf250f3a74d095ddf3c836d", size = 47571199, upload-time = "2026-02-16T10:13:42.504Z" }, + { url = "https://files.pythonhosted.org/packages/88/a3/d2c462d4ef313521eaf2eff04d204ac60775263f1fb08c374b543f79f610/pyarrow-23.0.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:46718a220d64677c93bc243af1d44b55998255427588e400677d7192671845c7", size = 48259435, upload-time = "2026-02-16T10:13:49.226Z" }, + { url = "https://files.pythonhosted.org/packages/cc/f1/11a544b8c3d38a759eb3fbb022039117fd633e9a7b19e4841cc3da091915/pyarrow-23.0.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a09f3876e87f48bc2f13583ab551f0379e5dfb83210391e68ace404181a20690", size = 50629149, upload-time = "2026-02-16T10:13:57.238Z" }, + { url = "https://files.pythonhosted.org/packages/50/f2/c0e76a0b451ffdf0cf788932e182758eb7558953f4f27f1aff8e2518b653/pyarrow-23.0.1-cp314-cp314t-win_amd64.whl", hash = "sha256:527e8d899f14bd15b740cd5a54ad56b7f98044955373a17179d5956ddb93d9ce", size = 28365807, upload-time = "2026-02-16T10:14:03.892Z" }, ] [[package]] @@ -1166,7 +1226,8 @@ dependencies = [ { name = "docutils" }, { name = "pygments" }, { name = "sphinx" }, - { name = "typing-extensions" }, + { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/00/20/bb50f9de3a6de69e6abd6b087b52fa2418a0418b19597601605f855ad044/pydata_sphinx_theme-0.16.1.tar.gz", hash = "sha256:a08b7f0b7f70387219dc659bff0893a7554d5eb39b59d3b8ef37b8401b7642d7", size = 2412693, upload-time = "2024-12-17T10:53:39.537Z" } wheels = [ @@ -1182,7 +1243,8 @@ dependencies = [ { name = "pyjwt", extra = ["crypto"] }, { name = "pynacl" }, { name = "requests" }, - { name = "typing-extensions" }, + { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "urllib3" }, ] sdist = { url = "https://files.pythonhosted.org/packages/16/ce/aa91d30040d9552c274e7ea8bd10a977600d508d579a4bb262b95eccf961/pygithub-2.5.0.tar.gz", hash = "sha256:e1613ac508a9be710920d26eb18b1905ebd9926aa49398e88151c1b526aad3cf", size = 3552804, upload-time = "2024-11-06T20:50:07.168Z" } @@ -1192,20 +1254,20 @@ wheels = [ [[package]] name = "pygments" -version = "2.19.1" +version = "2.20.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7c/2d/c3338d48ea6cc0feb8446d8e6937e1408088a72a39937982cc6111d17f84/pygments-2.19.1.tar.gz", hash = "sha256:61c16d2a8576dc0649d9f39e089b5f02bcd27fba10d8fb4dcc28173f7a45151f", size = 4968581, upload-time = "2025-01-06T17:26:30.443Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8a/0b/9fcc47d19c48b59121088dd6da2488a49d5f72dacf8262e2790a1d2c7d15/pygments-2.19.1-py3-none-any.whl", hash = "sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c", size = 1225293, upload-time = "2025-01-06T17:26:25.553Z" }, + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, ] [[package]] name = "pyjwt" -version = "2.10.1" +version = "2.12.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e7/46/bd74733ff231675599650d3e47f361794b22ef3e3770998dda30d3b63726/pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953", size = 87785, upload-time = "2024-11-28T03:43:29.933Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a8/10/e8192be5f38f3e8e7e046716de4cae33d56fd5ae08927a823bb916be36c1/pyjwt-2.12.0.tar.gz", hash = "sha256:2f62390b667cd8257de560b850bb5a883102a388829274147f1d724453f8fb02", size = 102511, upload-time = "2026-03-12T17:15:30.831Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/61/ad/689f02752eeec26aed679477e80e632ef1b682313be70793d798c1d5fc8f/PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb", size = 22997, upload-time = "2024-11-28T03:43:27.893Z" }, + { url = "https://files.pythonhosted.org/packages/15/70/70f895f404d363d291dcf62c12c85fdd47619ad9674ac0f53364d035925a/pyjwt-2.12.0-py3-none-any.whl", hash = "sha256:9bb459d1bdd0387967d287f5656bf7ec2b9a26645d1961628cda1764e087fd6e", size = 29700, upload-time = "2026-03-12T17:15:29.257Z" }, ] [package.optional-dependencies] @@ -1215,27 +1277,42 @@ crypto = [ [[package]] name = "pynacl" -version = "1.5.0" +version = "1.6.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi" }, + { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a7/22/27582568be639dfe22ddb3902225f91f2f17ceff88ce80e4db396c8986da/PyNaCl-1.5.0.tar.gz", hash = "sha256:8ac7448f09ab85811607bdd21ec2464495ac8b7c66d146bf545b0f08fb9220ba", size = 3392854, upload-time = "2022-01-07T22:05:41.134Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ce/75/0b8ede18506041c0bf23ac4d8e2971b4161cd6ce630b177d0a08eb0d8857/PyNaCl-1.5.0-cp36-abi3-macosx_10_10_universal2.whl", hash = "sha256:401002a4aaa07c9414132aaed7f6836ff98f59277a234704ff66878c2ee4a0d1", size = 349920, upload-time = "2022-01-07T22:05:49.156Z" }, - { url = "https://files.pythonhosted.org/packages/59/bb/fddf10acd09637327a97ef89d2a9d621328850a72f1fdc8c08bdf72e385f/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_24_aarch64.whl", hash = "sha256:52cb72a79269189d4e0dc537556f4740f7f0a9ec41c1322598799b0bdad4ef92", size = 601722, upload-time = "2022-01-07T22:05:50.989Z" }, - { url = "https://files.pythonhosted.org/packages/5d/70/87a065c37cca41a75f2ce113a5a2c2aa7533be648b184ade58971b5f7ccc/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a36d4a9dda1f19ce6e03c9a784a2921a4b726b02e1c736600ca9c22029474394", size = 680087, upload-time = "2022-01-07T22:05:52.539Z" }, - { url = "https://files.pythonhosted.org/packages/ee/87/f1bb6a595f14a327e8285b9eb54d41fef76c585a0edef0a45f6fc95de125/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_24_x86_64.whl", hash = "sha256:0c84947a22519e013607c9be43706dd42513f9e6ae5d39d3613ca1e142fba44d", size = 856678, upload-time = "2022-01-07T22:05:54.251Z" }, - { url = "https://files.pythonhosted.org/packages/66/28/ca86676b69bf9f90e710571b67450508484388bfce09acf8a46f0b8c785f/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:06b8f6fa7f5de8d5d2f7573fe8c863c051225a27b61e6860fd047b1775807858", size = 1133660, upload-time = "2022-01-07T22:05:56.056Z" }, - { url = "https://files.pythonhosted.org/packages/3d/85/c262db650e86812585e2bc59e497a8f59948a005325a11bbbc9ecd3fe26b/PyNaCl-1.5.0-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:a422368fc821589c228f4c49438a368831cb5bbc0eab5ebe1d7fac9dded6567b", size = 663824, upload-time = "2022-01-07T22:05:57.434Z" }, - { url = "https://files.pythonhosted.org/packages/fd/1a/cc308a884bd299b651f1633acb978e8596c71c33ca85e9dc9fa33a5399b9/PyNaCl-1.5.0-cp36-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:61f642bf2378713e2c2e1de73444a3778e5f0a38be6fee0fe532fe30060282ff", size = 1117912, upload-time = "2022-01-07T22:05:58.665Z" }, - { url = "https://files.pythonhosted.org/packages/25/2d/b7df6ddb0c2a33afdb358f8af6ea3b8c4d1196ca45497dd37a56f0c122be/PyNaCl-1.5.0-cp36-abi3-win32.whl", hash = "sha256:e46dae94e34b085175f8abb3b0aaa7da40767865ac82c928eeb9e57e1ea8a543", size = 204624, upload-time = "2022-01-07T22:06:00.085Z" }, - { url = "https://files.pythonhosted.org/packages/5e/22/d3db169895faaf3e2eda892f005f433a62db2decbcfbc2f61e6517adfa87/PyNaCl-1.5.0-cp36-abi3-win_amd64.whl", hash = "sha256:20f42270d27e1b6a29f54032090b972d97f0a1b0948cc52392041ef7831fee93", size = 212141, upload-time = "2022-01-07T22:06:01.861Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/d9/9a/4019b524b03a13438637b11538c82781a5eda427394380381af8f04f467a/pynacl-1.6.2.tar.gz", hash = "sha256:018494d6d696ae03c7e656e5e74cdfd8ea1326962cc401bcf018f1ed8436811c", size = 3511692, upload-time = "2026-01-01T17:48:10.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/79/0e3c34dc3c4671f67d251c07aa8eb100916f250ee470df230b0ab89551b4/pynacl-1.6.2-cp314-cp314t-macosx_10_10_universal2.whl", hash = "sha256:622d7b07cc5c02c666795792931b50c91f3ce3c2649762efb1ef0d5684c81594", size = 390064, upload-time = "2026-01-01T17:31:57.264Z" }, + { url = "https://files.pythonhosted.org/packages/eb/1c/23a26e931736e13b16483795c8a6b2f641bf6a3d5238c22b070a5112722c/pynacl-1.6.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d071c6a9a4c94d79eb665db4ce5cedc537faf74f2355e4d502591d850d3913c0", size = 809370, upload-time = "2026-01-01T17:31:59.198Z" }, + { url = "https://files.pythonhosted.org/packages/87/74/8d4b718f8a22aea9e8dcc8b95deb76d4aae380e2f5b570cc70b5fd0a852d/pynacl-1.6.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fe9847ca47d287af41e82be1dd5e23023d3c31a951da134121ab02e42ac218c9", size = 1408304, upload-time = "2026-01-01T17:32:01.162Z" }, + { url = "https://files.pythonhosted.org/packages/fd/73/be4fdd3a6a87fe8a4553380c2b47fbd1f7f58292eb820902f5c8ac7de7b0/pynacl-1.6.2-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:04316d1fc625d860b6c162fff704eb8426b1a8bcd3abacea11142cbd99a6b574", size = 844871, upload-time = "2026-01-01T17:32:02.824Z" }, + { url = "https://files.pythonhosted.org/packages/55/ad/6efc57ab75ee4422e96b5f2697d51bbcf6cdcc091e66310df91fbdc144a8/pynacl-1.6.2-cp314-cp314t-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:44081faff368d6c5553ccf55322ef2819abb40e25afaec7e740f159f74813634", size = 1446356, upload-time = "2026-01-01T17:32:04.452Z" }, + { url = "https://files.pythonhosted.org/packages/78/b7/928ee9c4779caa0a915844311ab9fb5f99585621c5d6e4574538a17dca07/pynacl-1.6.2-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:a9f9932d8d2811ce1a8ffa79dcbdf3970e7355b5c8eb0c1a881a57e7f7d96e88", size = 826814, upload-time = "2026-01-01T17:32:06.078Z" }, + { url = "https://files.pythonhosted.org/packages/f7/a9/1bdba746a2be20f8809fee75c10e3159d75864ef69c6b0dd168fc60e485d/pynacl-1.6.2-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:bc4a36b28dd72fb4845e5d8f9760610588a96d5a51f01d84d8c6ff9849968c14", size = 1411742, upload-time = "2026-01-01T17:32:07.651Z" }, + { url = "https://files.pythonhosted.org/packages/f3/2f/5e7ea8d85f9f3ea5b6b87db1d8388daa3587eed181bdeb0306816fdbbe79/pynacl-1.6.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3bffb6d0f6becacb6526f8f42adfb5efb26337056ee0831fb9a7044d1a964444", size = 801714, upload-time = "2026-01-01T17:32:09.558Z" }, + { url = "https://files.pythonhosted.org/packages/06/ea/43fe2f7eab5f200e40fb10d305bf6f87ea31b3bbc83443eac37cd34a9e1e/pynacl-1.6.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:2fef529ef3ee487ad8113d287a593fa26f48ee3620d92ecc6f1d09ea38e0709b", size = 1372257, upload-time = "2026-01-01T17:32:11.026Z" }, + { url = "https://files.pythonhosted.org/packages/4d/54/c9ea116412788629b1347e415f72195c25eb2f3809b2d3e7b25f5c79f13a/pynacl-1.6.2-cp314-cp314t-win32.whl", hash = "sha256:a84bf1c20339d06dc0c85d9aea9637a24f718f375d861b2668b2f9f96fa51145", size = 231319, upload-time = "2026-01-01T17:32:12.46Z" }, + { url = "https://files.pythonhosted.org/packages/ce/04/64e9d76646abac2dccf904fccba352a86e7d172647557f35b9fe2a5ee4a1/pynacl-1.6.2-cp314-cp314t-win_amd64.whl", hash = "sha256:320ef68a41c87547c91a8b58903c9caa641ab01e8512ce291085b5fe2fcb7590", size = 244044, upload-time = "2026-01-01T17:32:13.781Z" }, + { url = "https://files.pythonhosted.org/packages/33/33/7873dc161c6a06f43cda13dec67b6fe152cb2f982581151956fa5e5cdb47/pynacl-1.6.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d29bfe37e20e015a7d8b23cfc8bd6aa7909c92a1b8f41ee416bbb3e79ef182b2", size = 188740, upload-time = "2026-01-01T17:32:15.083Z" }, + { url = "https://files.pythonhosted.org/packages/be/7b/4845bbf88e94586ec47a432da4e9107e3fc3ce37eb412b1398630a37f7dd/pynacl-1.6.2-cp38-abi3-macosx_10_10_universal2.whl", hash = "sha256:c949ea47e4206af7c8f604b8278093b674f7c79ed0d4719cc836902bf4517465", size = 388458, upload-time = "2026-01-01T17:32:16.829Z" }, + { url = "https://files.pythonhosted.org/packages/1e/b4/e927e0653ba63b02a4ca5b4d852a8d1d678afbf69b3dbf9c4d0785ac905c/pynacl-1.6.2-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8845c0631c0be43abdd865511c41eab235e0be69c81dc66a50911594198679b0", size = 800020, upload-time = "2026-01-01T17:32:18.34Z" }, + { url = "https://files.pythonhosted.org/packages/7f/81/d60984052df5c97b1d24365bc1e30024379b42c4edcd79d2436b1b9806f2/pynacl-1.6.2-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:22de65bb9010a725b0dac248f353bb072969c94fa8d6b1f34b87d7953cf7bbe4", size = 1399174, upload-time = "2026-01-01T17:32:20.239Z" }, + { url = "https://files.pythonhosted.org/packages/68/f7/322f2f9915c4ef27d140101dd0ed26b479f7e6f5f183590fd32dfc48c4d3/pynacl-1.6.2-cp38-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:46065496ab748469cdd999246d17e301b2c24ae2fdf739132e580a0e94c94a87", size = 835085, upload-time = "2026-01-01T17:32:22.24Z" }, + { url = "https://files.pythonhosted.org/packages/3e/d0/f301f83ac8dbe53442c5a43f6a39016f94f754d7a9815a875b65e218a307/pynacl-1.6.2-cp38-abi3-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8a66d6fb6ae7661c58995f9c6435bda2b1e68b54b598a6a10247bfcdadac996c", size = 1437614, upload-time = "2026-01-01T17:32:23.766Z" }, + { url = "https://files.pythonhosted.org/packages/c4/58/fc6e649762b029315325ace1a8c6be66125e42f67416d3dbd47b69563d61/pynacl-1.6.2-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:26bfcd00dcf2cf160f122186af731ae30ab120c18e8375684ec2670dccd28130", size = 818251, upload-time = "2026-01-01T17:32:25.69Z" }, + { url = "https://files.pythonhosted.org/packages/c9/a8/b917096b1accc9acd878819a49d3d84875731a41eb665f6ebc826b1af99e/pynacl-1.6.2-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:c8a231e36ec2cab018c4ad4358c386e36eede0319a0c41fed24f840b1dac59f6", size = 1402859, upload-time = "2026-01-01T17:32:27.215Z" }, + { url = "https://files.pythonhosted.org/packages/85/42/fe60b5f4473e12c72f977548e4028156f4d340b884c635ec6b063fe7e9a5/pynacl-1.6.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:68be3a09455743ff9505491220b64440ced8973fe930f270c8e07ccfa25b1f9e", size = 791926, upload-time = "2026-01-01T17:32:29.314Z" }, + { url = "https://files.pythonhosted.org/packages/fa/f9/e40e318c604259301cc091a2a63f237d9e7b424c4851cafaea4ea7c4834e/pynacl-1.6.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8b097553b380236d51ed11356c953bf8ce36a29a3e596e934ecabe76c985a577", size = 1363101, upload-time = "2026-01-01T17:32:31.263Z" }, + { url = "https://files.pythonhosted.org/packages/48/47/e761c254f410c023a469284a9bc210933e18588ca87706ae93002c05114c/pynacl-1.6.2-cp38-abi3-win32.whl", hash = "sha256:5811c72b473b2f38f7e2a3dc4f8642e3a3e9b5e7317266e4ced1fba85cae41aa", size = 227421, upload-time = "2026-01-01T17:32:33.076Z" }, + { url = "https://files.pythonhosted.org/packages/41/ad/334600e8cacc7d86587fe5f565480fde569dfb487389c8e1be56ac21d8ac/pynacl-1.6.2-cp38-abi3-win_amd64.whl", hash = "sha256:62985f233210dee6548c223301b6c25440852e13d59a8b81490203c3227c5ba0", size = 239754, upload-time = "2026-01-01T17:32:34.557Z" }, + { url = "https://files.pythonhosted.org/packages/29/7d/5945b5af29534641820d3bd7b00962abbbdfee84ec7e19f0d5b3175f9a31/pynacl-1.6.2-cp38-abi3-win_arm64.whl", hash = "sha256:834a43af110f743a754448463e8fd61259cd4ab5bbedcf70f9dabad1d28a394c", size = 184801, upload-time = "2026-01-01T17:32:36.309Z" }, ] [[package]] name = "pytest" -version = "8.3.4" +version = "9.0.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, @@ -1243,23 +1320,27 @@ dependencies = [ { name = "iniconfig" }, { name = "packaging" }, { name = "pluggy" }, + { name = "pygments" }, { name = "tomli", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/05/35/30e0d83068951d90a01852cb1cef56e5d8a09d20c7f511634cc2f7e0372a/pytest-8.3.4.tar.gz", hash = "sha256:965370d062bce11e73868e0335abac31b4d3de0e82f4007408d242b4f8610761", size = 1445919, upload-time = "2024-12-01T12:54:25.98Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/11/92/76a1c94d3afee238333bc0a42b82935dd8f9cf8ce9e336ff87ee14d9e1cf/pytest-8.3.4-py3-none-any.whl", hash = "sha256:50e16d954148559c9a74109af1eaf0c945ba2d8f30f0a3d3335edde19788b6f6", size = 343083, upload-time = "2024-12-01T12:54:19.735Z" }, + { url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" }, ] [[package]] name = "pytest-asyncio" -version = "0.25.3" +version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "backports-asyncio-runner", marker = "python_full_version < '3.11'" }, { name = "pytest" }, + { name = "typing-extensions", version = "4.12.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "typing-extensions", version = "4.15.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f2/a8/ecbc8ede70921dd2f544ab1cadd3ff3bf842af27f87bbdea774c7baa1d38/pytest_asyncio-0.25.3.tar.gz", hash = "sha256:fc1da2cf9f125ada7e710b4ddad05518d4cee187ae9412e9ac9271003497f07a", size = 54239, upload-time = "2025-01-28T18:37:58.729Z" } +sdist = { url = "https://files.pythonhosted.org/packages/43/7c/d36d04db312ecf4298932ef77e6e4a9e8ad017906e24e34f0b0c361a2473/pytest_asyncio-1.4.0.tar.gz", hash = "sha256:c6c0d2259945122819f171a32ecea2c349ead889ee28176caaf492143424be42", size = 58514, upload-time = "2026-05-26T09:56:04.083Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/67/17/3493c5624e48fd97156ebaec380dcaafee9506d7e2c46218ceebbb57d7de/pytest_asyncio-0.25.3-py3-none-any.whl", hash = "sha256:9e89518e0f9bd08928f97a3482fdc4e244df17529460bc038291ccaf8f85c7c3", size = 19467, upload-time = "2025-01-28T18:37:56.798Z" }, + { url = "https://files.pythonhosted.org/packages/03/e2/08a497ef684b88559c9cc5f4ad53a37e7b99e727094a86d6ea32536d5d3c/pytest_asyncio-1.4.0-py3-none-any.whl", hash = "sha256:933ca923a23075a87fb7070c0ec272a6848489824d887c85c812670932835aa1", size = 16930, upload-time = "2026-05-26T09:56:02.576Z" }, ] [[package]] @@ -1361,7 +1442,7 @@ wheels = [ [[package]] name = "requests" -version = "2.32.3" +version = "2.33.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "certifi" }, @@ -1369,9 +1450,9 @@ dependencies = [ { name = "idna" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/63/70/2bf7780ad2d390a8d301ad0b550f1581eadbd9a20f896afe06353c2a2913/requests-2.32.3.tar.gz", hash = "sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760", size = 131218, upload-time = "2024-05-29T15:37:49.536Z" } +sdist = { url = "https://files.pythonhosted.org/packages/34/64/8860370b167a9721e8956ae116825caff829224fbca0ca6e7bf8ddef8430/requests-2.33.0.tar.gz", hash = "sha256:c7ebc5e8b0f21837386ad0e1c8fe8b829fa5f544d8df3b2253bff14ef29d7652", size = 134232, upload-time = "2026-03-25T15:10:41.586Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/9b/335f9764261e915ed497fcdeb11df5dfd6f7bf257d4a6a2a686d80da4d54/requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6", size = 64928, upload-time = "2024-05-29T15:37:47.027Z" }, + { url = "https://files.pythonhosted.org/packages/56/5d/c814546c2333ceea4ba42262d8c4d55763003e767fa169adc693bd524478/requests-2.33.0-py3-none-any.whl", hash = "sha256:3324635456fa185245e24865e810cecec7b4caf933d7eb133dcde67d48cee69b", size = 65017, upload-time = "2026-03-25T15:10:40.382Z" }, ] [[package]] @@ -1639,11 +1720,28 @@ wheels = [ name = "typing-extensions" version = "4.12.2" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version >= '3.12' and python_full_version < '3.14'", + "python_full_version == '3.11.*'", +] sdist = { url = "https://files.pythonhosted.org/packages/df/db/f35a00659bc03fec321ba8bce9420de607a1d37f8342eee1863174c69557/typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8", size = 85321, upload-time = "2024-06-07T18:52:15.995Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/26/9f/ad63fc0248c5379346306f8668cda6e2e2e9c95e01216d2b8ffd9ff037d0/typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d", size = 37438, upload-time = "2024-06-07T18:52:13.582Z" }, ] +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] + [[package]] name = "tzdata" version = "2024.2" @@ -1655,11 +1753,11 @@ wheels = [ [[package]] name = "urllib3" -version = "2.3.0" +version = "2.7.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/aa/63/e53da845320b757bf29ef6a9062f5c669fe997973f966045cb019c3f4b66/urllib3-2.3.0.tar.gz", hash = "sha256:f8c5449b3cf0861679ce7e0503c7b44b5ec981bec0d1d3795a07f1ba96f0204d", size = 307268, upload-time = "2024-12-22T07:47:30.032Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/19/4ec628951a74043532ca2cf5d97b7b14863931476d117c471e8e2b1eb39f/urllib3-2.3.0-py3-none-any.whl", hash = "sha256:1cee9ad369867bfdbbb48b7dd50374c0967a0bb7710050facf0dd6911440e3df", size = 128369, upload-time = "2024-12-22T07:47:28.074Z" }, + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] [[package]] From 3256138ecc57de2666df704909a371919ff6fbfd Mon Sep 17 00:00:00 2001 From: kosiew Date: Sat, 13 Jun 2026 23:40:08 +0800 Subject: [PATCH 59/83] Deprecate `Expr` temporal part arguments in date extraction and truncation functions (#1587) Added a shared helper to emit a DeprecationWarning when an Expr is passed to a literal control argument. Updated the implementations of: date_part datepart extract date_trunc datetrunc to warn when part is provided as an Expr. Refactored alias functions (datepart, extract, and datetrunc) to delegate through internal helper functions so warnings reference the user-facing function name correctly and avoid duplicate warning behavior. Updated existing temporal function tests to use the preferred string-literal form for part. Added targeted tests verifying: Expr inputs emit DeprecationWarning for date_part, datepart, and extract. Expr inputs emit DeprecationWarning for date_trunc and datetrunc. String literal inputs continue to work without emitting deprecation warnings. --- .../common-operations/functions.rst | 4 +- python/datafusion/functions.py | 29 +++++++++-- python/tests/test_functions.py | 50 ++++++++++++++++--- 3 files changed, 71 insertions(+), 12 deletions(-) diff --git a/docs/source/user-guide/common-operations/functions.rst b/docs/source/user-guide/common-operations/functions.rst index ccb47a4e7..ed5a991d8 100644 --- a/docs/source/user-guide/common-operations/functions.rst +++ b/docs/source/user-guide/common-operations/functions.rst @@ -77,8 +77,8 @@ Extracting parts of a date using :py:func:`~datafusion.functions.date_part` (ali .. ipython:: python df.select( - f.date_part(literal("month"), f.to_timestamp(col('"Total"'))).alias("month"), - f.extract(literal("day"), f.to_timestamp(col('"Total"'))).alias("day") + f.date_part("month", f.to_timestamp(col('"Total"'))).alias("month"), + f.extract("day", f.to_timestamp(col('"Total"'))).alias("day") ) String diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index c8f07497d..9158a7146 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -39,6 +39,7 @@ from __future__ import annotations import inspect +import warnings from typing import TYPE_CHECKING, Any import pyarrow as pa @@ -60,6 +61,16 @@ sort_or_default, ) + +def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: + warnings.warn( + f"Passing Expr for {function_name}() argument {arg_name!r} is deprecated; " + "pass a Python literal instead.", + DeprecationWarning, + stacklevel=4, + ) + + __all__ = [ "abs", "acos", @@ -2575,7 +2586,7 @@ def datepart(part: Expr | str, date: Expr) -> Expr: See Also: This is an alias for :py:func:`date_part`. """ - return date_part(part, date) + return _date_part(part, date, "datepart") def date_part(part: Expr | str, date: Expr) -> Expr: @@ -2595,6 +2606,12 @@ def date_part(part: Expr | str, date: Expr) -> Expr: >>> result.collect_column("y")[0].as_py() 2021 """ + return _date_part(part, date, "date_part") + + +def _date_part(part: Expr | str, date: Expr, function_name: str) -> Expr: + if isinstance(part, Expr): + _warn_expr_for_literal_arg(function_name, "part") part = coerce_to_expr(part) return Expr(f.date_part(part.expr, date.expr)) @@ -2605,7 +2622,7 @@ def extract(part: Expr | str, date: Expr) -> Expr: See Also: This is an alias for :py:func:`date_part`. """ - return date_part(part, date) + return _date_part(part, date, "extract") def date_trunc(part: Expr | str, date: Expr) -> Expr: @@ -2626,6 +2643,12 @@ def date_trunc(part: Expr | str, date: Expr) -> Expr: >>> str(result.collect_column("t")[0].as_py()) '2021-07-01 00:00:00' """ + return _date_trunc(part, date, "date_trunc") + + +def _date_trunc(part: Expr | str, date: Expr, function_name: str) -> Expr: + if isinstance(part, Expr): + _warn_expr_for_literal_arg(function_name, "part") part = coerce_to_expr(part) return Expr(f.date_trunc(part.expr, date.expr)) @@ -2636,7 +2659,7 @@ def datetrunc(part: Expr | str, date: Expr) -> Expr: See Also: This is an alias for :py:func:`date_trunc`. """ - return date_trunc(part, date) + return _date_trunc(part, date, "datetrunc") def date_bin(stride: Expr, source: Expr, origin: Expr) -> Expr: diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index 55d9c8ee8..d2abe4741 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import math +import warnings from datetime import date, datetime, time, timezone import numpy as np @@ -1086,10 +1087,10 @@ def test_hash_functions(df): def test_temporal_functions(df): df = df.select( - f.date_part(literal("month"), column("d")), - f.datepart(literal("year"), column("d")), - f.date_trunc(literal("month"), column("d")), - f.datetrunc(literal("day"), column("d")), + f.date_part("month", column("d")), + f.datepart("year", column("d")), + f.date_trunc("month", column("d")), + f.datetrunc("day", column("d")), f.date_bin( literal("15 minutes").cast(pa.string()), column("d"), @@ -1100,7 +1101,7 @@ def test_temporal_functions(df): f.to_timestamp_seconds(literal("2023-09-07 05:06:14.523952")), f.to_timestamp_millis(literal("2023-09-07 05:06:14.523952")), f.to_timestamp_micros(literal("2023-09-07 05:06:14.523952")), - f.extract(literal("day"), column("d")), + f.extract("day", column("d")), f.to_timestamp( literal("2023-09-07 05:06:14.523952000"), literal("%Y-%m-%d %H:%M:%S.%f") ), @@ -2160,16 +2161,51 @@ def test_date_part_native_str(self): ctx = SessionContext() df = ctx.from_pydict({"a": ["2021-07-15T00:00:00"]}) df = df.select(f.to_timestamp(column("a")).alias("a")) - result = df.select(f.date_part("year", column("a")).alias("y")).collect() + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + result = df.select(f.date_part("year", column("a")).alias("y")).collect() assert result[0].column(0)[0].as_py() == 2021 + @pytest.mark.parametrize( + ("func", "name"), + [ + pytest.param(f.date_part, "date_part", id="date_part"), + pytest.param(f.datepart, "datepart", id="datepart"), + pytest.param(f.extract, "extract", id="extract"), + ], + ) + def test_date_part_expr_part_warns_deprecated(self, func, name): + with pytest.warns( + DeprecationWarning, + match=rf"Passing Expr for {name}\(\) argument 'part' is deprecated", + ): + expr = func(literal("year"), column("a")) + assert expr is not None + def test_date_trunc_native_str(self): ctx = SessionContext() df = ctx.from_pydict({"a": ["2021-07-15T12:34:56"]}) df = df.select(f.to_timestamp(column("a")).alias("a")) - result = df.select(f.date_trunc("month", column("a")).alias("t")).collect() + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + result = df.select(f.date_trunc("month", column("a")).alias("t")).collect() assert str(result[0].column(0)[0].as_py()) == "2021-07-01 00:00:00" + @pytest.mark.parametrize( + ("func", "name"), + [ + pytest.param(f.date_trunc, "date_trunc", id="date_trunc"), + pytest.param(f.datetrunc, "datetrunc", id="datetrunc"), + ], + ) + def test_date_trunc_expr_part_warns_deprecated(self, func, name): + with pytest.warns( + DeprecationWarning, + match=rf"Passing Expr for {name}\(\) argument 'part' is deprecated", + ): + expr = func(literal("month"), column("a")) + assert expr is not None + def test_left_native_int(self): ctx = SessionContext() df = ctx.from_pydict({"a": ["the cat"]}) From fd15b03ecc227957b752429bf5b7c56c0f139f1e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 15 Jun 2026 14:01:13 -0400 Subject: [PATCH 60/83] feat: expose array_compact, array_normalize, cosine_distance, inner_product (#1567) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: expose array_compact, array_normalize, cosine_distance, inner_product Adds Python bindings for four scalar functions from datafusion::functions_nested::expr_fn that were not previously surfaced: - array_compact / list_compact: drop NULLs from an array. - array_normalize / list_normalize: L2-normalize a numeric array. - cosine_distance: 1 - cosine_similarity(a, b). - inner_product: dot product of two numeric arrays. Implementation routes each through the existing array_fn! macro in crates/core/src/functions.rs, mirroring the other functions_nested wrappers. Python wrappers in python/datafusion/functions.py follow the established pattern with doctest examples; list_* aliases use the one-line + See Also form per project convention. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: clarify array_normalize and cosine_distance docstrings Expand both docstrings with plain-English definitions, worked examples, ranges, use cases, and behavior on edge cases (zero vector → NULL, length-mismatched inputs fail). Adds a zero-vector example to array_normalize and an orthogonal-vector example to cosine_distance. Updates the list_normalize alias summary to match. Co-Authored-By: Claude * test: add alias-equivalence and length-mismatch tests for array fns Pin the contracts the doctests don't cover: list_compact/list_normalize must produce the same output as their array_* primaries, and cosine_distance/inner_product must reject length-mismatched inputs at execution time. * feat: expose dot_product alias for inner_product Match upstream DataFusion SQL alias surface (inner_product UDF registers `dot_product` in its alias list). Also expand `inner_product` docstring with NULL/length-mismatch behavior to match peer distance fns added in this PR. Co-Authored-By: Claude Opus 4.7 * test: fold dot_product alias check into parametrized test Generalize test_array_function_aliases to accept multi-column data so the dot_product/inner_product alias case fits, dropping the standalone test_dot_product_alias_matches_inner_product. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/functions.rs | 8 ++ python/datafusion/functions.py | 183 +++++++++++++++++++++++++++++++++ python/tests/test_functions.py | 33 ++++++ 3 files changed, 224 insertions(+) diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index 395d5ebfd..a56873c58 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -654,6 +654,10 @@ array_fn!(array_replace, array from to); array_fn!(array_replace_n, array from to max); array_fn!(array_replace_all, array from to); array_fn!(array_sort, array desc null_first); +array_fn!(array_compact, array); +array_fn!(array_normalize, array); +array_fn!(cosine_distance, array1 array2); +array_fn!(inner_product, array1 array2); array_fn!(array_intersect, first_array second_array); array_fn!(array_union, array1 array2); array_fn!(array_except, first_array second_array); @@ -1133,6 +1137,10 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(array_cat))?; m.add_wrapped(wrap_pyfunction!(array_dims))?; m.add_wrapped(wrap_pyfunction!(array_distinct))?; + m.add_wrapped(wrap_pyfunction!(array_compact))?; + m.add_wrapped(wrap_pyfunction!(array_normalize))?; + m.add_wrapped(wrap_pyfunction!(cosine_distance))?; + m.add_wrapped(wrap_pyfunction!(inner_product))?; m.add_wrapped(wrap_pyfunction!(array_element))?; m.add_wrapped(wrap_pyfunction!(array_empty))?; m.add_wrapped(wrap_pyfunction!(array_length))?; diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 9158a7146..4e08fa1d9 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -87,6 +87,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "array_any_value", "array_append", "array_cat", + "array_compact", "array_concat", "array_contains", "array_dims", @@ -107,6 +108,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "array_max", "array_min", "array_ndims", + "array_normalize", "array_pop_back", "array_pop_front", "array_position", @@ -162,6 +164,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "corr", "cos", "cosh", + "cosine_distance", "cot", "count", "count_star", @@ -182,6 +185,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "degrees", "dense_rank", "digest", + "dot_product", "element_at", "empty", "encode", @@ -203,6 +207,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "ifnull", "in_list", "initcap", + "inner_product", "instr", "isnan", "iszero", @@ -220,6 +225,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "list_any_value", "list_append", "list_cat", + "list_compact", "list_concat", "list_contains", "list_dims", @@ -240,6 +246,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "list_max", "list_min", "list_ndims", + "list_normalize", "list_overlap", "list_pop_back", "list_pop_front", @@ -3227,6 +3234,164 @@ def array_distinct(array: Expr) -> Expr: return Expr(f.array_distinct(array.expr)) +def array_compact(array: Expr) -> Expr: + """Removes NULL values from the array. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, None, 2, None, 3]]}) + >>> result = df.select( + ... dfn.functions.array_compact(dfn.col("a")).alias("result") + ... ) + >>> result.collect_column("result")[0].as_py() + [1, 2, 3] + """ + return Expr(f.array_compact(array.expr)) + + +def array_normalize(array: Expr) -> Expr: + """Scales a numeric array so it has Euclidean length 1. + + Treats the array as a vector and divides every element by the vector's + Euclidean (L2) norm — the square root of the sum of the squared + elements. The returned array points in the same direction as the input + but has a magnitude of 1, which makes it suitable for cosine-similarity + comparisons and other operations that expect unit vectors. + + For the input ``[3.0, 4.0]`` the L2 norm is ``sqrt(3**2 + 4**2) = 5``, + so each element is divided by 5 to produce ``[0.6, 0.8]``. + + Normalizing the zero vector is undefined (it would divide by zero), so + the function returns NULL for an all-zero input. NULL is also returned + if any element of the input array is NULL. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[3.0, 4.0]]}) + >>> result = df.select( + ... dfn.functions.array_normalize(dfn.col("a")).alias("result") + ... ) + >>> result.collect_column("result")[0].as_py() + [0.6, 0.8] + + The zero vector has no direction to preserve, so the result is NULL: + + >>> df_zero = ctx.from_pydict({"a": [[0.0, 0.0]]}) + >>> result = df_zero.select( + ... dfn.functions.array_normalize(dfn.col("a")).alias("result") + ... ) + >>> result.collect_column("result")[0].as_py() is None + True + """ + return Expr(f.array_normalize(array.expr)) + + +def cosine_distance(array1: Expr, array2: Expr) -> Expr: + """Measures how much two numeric arrays differ in direction. + + Treats each input as a vector and compares the angle between them, + ignoring their magnitudes. The result is ``1 - cosine_similarity``, + where cosine similarity is the dot product of the two vectors divided + by the product of their Euclidean (L2) norms. + + The returned value ranges from 0 to 2: + + * ``0`` — vectors point in the same direction (any positive scaling + of one yields the other). + * ``1`` — vectors are orthogonal (no shared direction). + * ``2`` — vectors point in exactly opposite directions. + + This is the standard distance metric for comparing embedding vectors + (text, image, audio) where direction carries the meaning and overall + magnitude does not. + + Both arrays must have the same length; otherwise execution fails. If + either input is the zero vector the cosine is undefined and the + function returns NULL. + + Examples: + Identical vectors have distance ``0``: + + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict( + ... {"a": [[1.0, 2.0, 3.0]], "b": [[1.0, 2.0, 3.0]]} + ... ) + >>> result = df.select( + ... dfn.functions.cosine_distance( + ... dfn.col("a"), dfn.col("b") + ... ).alias("result") + ... ) + >>> result.collect_column("result")[0].as_py() + 0.0 + + Orthogonal vectors have distance ``1``: + + >>> df_orth = ctx.from_pydict( + ... {"a": [[1.0, 0.0]], "b": [[0.0, 1.0]]} + ... ) + >>> result = df_orth.select( + ... dfn.functions.cosine_distance( + ... dfn.col("a"), dfn.col("b") + ... ).alias("result") + ... ) + >>> result.collect_column("result")[0].as_py() + 1.0 + """ + return Expr(f.cosine_distance(array1.expr, array2.expr)) + + +def inner_product(array1: Expr, array2: Expr) -> Expr: + """Returns the inner (dot) product of two numeric arrays. + + Treats each input as a vector and returns the sum of the element-wise + products: ``sum(array1[i] * array2[i])``. For ``[1, 2, 3]`` and + ``[4, 5, 6]`` the result is ``1*4 + 2*5 + 3*6 = 32``. + + Also available as :py:func:`dot_product` (and as ``dot_product`` in + raw SQL). + + Both arrays must have the same length; otherwise execution fails. NULL + is returned when either input array is NULL or when any element of + either array is NULL. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict( + ... {"a": [[1.0, 2.0, 3.0]], "b": [[4.0, 5.0, 6.0]]} + ... ) + >>> result = df.select( + ... dfn.functions.inner_product( + ... dfn.col("a"), dfn.col("b") + ... ).alias("result") + ... ) + >>> result.collect_column("result")[0].as_py() + 32.0 + + NULL elements propagate to NULL output: + + >>> df_null = ctx.from_pydict( + ... {"a": [[1.0, None, 3.0]], "b": [[4.0, 5.0, 6.0]]} + ... ) + >>> result = df_null.select( + ... dfn.functions.inner_product( + ... dfn.col("a"), dfn.col("b") + ... ).alias("result") + ... ) + >>> result.collect_column("result")[0].as_py() is None + True + """ + return Expr(f.inner_product(array1.expr, array2.expr)) + + +def dot_product(array1: Expr, array2: Expr) -> Expr: + """Returns the inner (dot) product of two numeric arrays. + + See Also: + This is an alias for :py:func:`inner_product`. + """ + return inner_product(array1, array2) + + def list_cat(*args: Expr) -> Expr: """Concatenates the input arrays. @@ -3254,6 +3419,24 @@ def list_distinct(array: Expr) -> Expr: return array_distinct(array) +def list_compact(array: Expr) -> Expr: + """Removes NULL values from the array. + + See Also: + This is an alias for :py:func:`array_compact`. + """ + return array_compact(array) + + +def list_normalize(array: Expr) -> Expr: + """Scales a numeric array so it has Euclidean length 1. + + See Also: + This is an alias for :py:func:`array_normalize`. + """ + return array_normalize(array) + + def list_dims(array: Expr) -> Expr: """Returns an array of the array's dimensions. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index d2abe4741..b02cdf6db 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -718,6 +718,39 @@ def test_array_function_obj_tests(stmt, py_expr): assert a == b +@pytest.mark.parametrize( + ("alias_fn", "primary_fn", "data"), + [ + (f.list_compact, f.array_compact, {"a": [[1.0, None, 2.0, None, 3.0]]}), + (f.list_normalize, f.array_normalize, {"a": [[3.0, 4.0]]}), + ( + f.dot_product, + f.inner_product, + {"a": [[1.0, 2.0, 3.0]], "b": [[4.0, 5.0, 6.0]]}, + ), + ], +) +def test_array_function_aliases(alias_fn, primary_fn, data): + """Alias helpers should be exact aliases for their primary counterparts.""" + ctx = SessionContext() + df = ctx.from_pydict(data) + cols = [column(name) for name in data] + alias_result = df.select(alias_fn(*cols).alias("r")).collect() + primary_result = df.select(primary_fn(*cols).alias("r")).collect() + assert ( + alias_result[0].column(0).to_pylist() == primary_result[0].column(0).to_pylist() + ) + + +@pytest.mark.parametrize("fn", [f.cosine_distance, f.inner_product, f.dot_product]) +def test_array_distance_length_mismatch_raises(fn): + """Length-mismatched inputs to vector distance fns should raise at execute.""" + ctx = SessionContext() + df = ctx.from_pydict({"a": [[1.0, 2.0]], "b": [[1.0, 2.0, 3.0]]}) + with pytest.raises(Exception, match="same length"): + df.select(fn(column("a"), column("b")).alias("r")).collect() + + @pytest.mark.parametrize( ("args", "expected"), [ From c0ac93b68fff606aed570aed09d72d749ba79bae Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Mon, 15 Jun 2026 16:23:05 -0400 Subject: [PATCH 61/83] feat: expose arrow_field, arrow_try_cast, cast_to_type, with_metadata (#1568) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: expose arrow_field, arrow_try_cast, cast_to_type, with_metadata Adds Python bindings for five scalar functions from datafusion::functions::expr_fn that were not previously surfaced: - arrow_field: returns a struct describing an expression's Arrow field (name, data_type, nullable, metadata). - arrow_try_cast: like arrow_cast but yields NULL on cast failure. - cast_to_type / try_cast_to_type: casts a value to the type of a reference expression. These are exposed as a single Python entry point cast_to_type(value, type_ref, *, try_cast=False); the kwarg switches between the strict and try variants. - with_metadata: attach Arrow field metadata; the inverse of arrow_metadata. Accepts a dict[str, str] for ergonomics. Updates skills/datafusion_python/SKILL.md to list the new functions and documents the cast_to_type kwarg behavior. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: collapse try_cast_to_type into cast_to_type kwarg The previous commit exposed cast_to_type and try_cast_to_type as two separate pyo3 bindings and unified them in the Python wrapper via a try_cast kwarg. That left try_cast_to_type in datafusion._internal without a matching public Python name, breaking test_datafusion_missing_exports. Move the dispatch into the rust binding: cast_to_type now takes a try_cast kwarg and selects between functions::expr_fn::cast_to_type and try_cast_to_type internally. Only one pyo3 binding is registered, so the wrapper-coverage check passes and the Python entrypoint is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) * feat: accept pyarrow DataType in arrow_try_cast Mirrors arrow_cast: arrow_try_cast now accepts `pa.DataType` in addition to `str` and `Expr`. Adds `Expr.try_cast(pa.DataType)` PyO3 binding for the pyarrow-type routing path. Co-Authored-By: Claude Opus 4.7 (1M context) * fix: guard with_metadata against empty dict and empty keys Empty `metadata` dict now returns the input expression unchanged (previously bubbled an opaque DataFusion error about minimum arg count). Empty keys raise `ValueError` to match the docstring contract. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: assert full struct shape in arrow_field doctest Previous doctest set metadata on the input field but only checked the name — the metadata setup was dead. Now the example asserts the full returned struct (name, data_type, nullable, metadata) so the demo shows what the function actually produces. Co-Authored-By: Claude Opus 4.7 (1M context) * test: add unit tests for arrow_try_cast, arrow_field, cast_to_type, with_metadata Mirrors the existing test_arrow_cast pattern. Covers: - arrow_try_cast: string-syntax, pa.DataType, and null-on-failure paths - arrow_field: full returned struct shape (name, data_type, nullable, metadata) - cast_to_type: type-from-expr happy path and try_cast=True null behavior - with_metadata: round-trip through arrow_metadata, empty-dict no-op, and empty-key ValueError Co-Authored-By: Claude Opus 4.7 (1M context) * test: parameterize arrow cast / try_cast tests Folds the previous four cast tests (arrow_cast + arrow_try_cast × str + pyarrow target type) into a single parameterized test that runs both functions across all five target-type variants. Collapses the two cast_to_type tests (happy path + try_cast=True) into one parameterized test, and parameterizes arrow_try_cast null-on-failure over both target-type syntaxes. 7 test functions, 19 cases — net less code, same coverage. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: point cast_to_type at arrow_cast for static target types Adds a one-line cross-reference so users with a known target type reach for arrow_cast / arrow_try_cast instead of building a sentinel expression to feed cast_to_type. Co-Authored-By: Claude Opus 4.7 (1M context) * refactor: split cast_to_type into cast_to_type and try_cast_to_type Replace the try_cast bool flag with separate cast_to_type and try_cast_to_type functions, matching upstream DataFusion and the arrow_cast / arrow_try_cast pair. Also drop the redundant data_type parametrization on test_arrow_try_cast_null_on_failure, since the str-vs-pyarrow distinction is already covered by test_arrow_cast_variants. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- crates/core/src/expr.rs | 5 ++ crates/core/src/functions.rs | 10 +++ python/datafusion/expr.py | 22 +++++ python/datafusion/functions.py | 144 ++++++++++++++++++++++++++++++ python/tests/test_functions.py | 96 ++++++++++++++++---- skills/datafusion_python/SKILL.md | 7 +- 6 files changed, 265 insertions(+), 19 deletions(-) diff --git a/crates/core/src/expr.rs b/crates/core/src/expr.rs index eac571a11..432c4cd23 100644 --- a/crates/core/src/expr.rs +++ b/crates/core/src/expr.rs @@ -358,6 +358,11 @@ impl PyExpr { expr.into() } + pub fn try_cast(&self, to: PyArrowType) -> PyExpr { + let expr = Expr::TryCast(TryCast::new(Box::new(self.expr.clone()), to.0)); + expr.into() + } + #[pyo3(signature = (low, high, negated=false))] pub fn between(&self, low: PyExpr, high: PyExpr, negated: bool) -> PyExpr { let expr = Expr::Between(Between::new( diff --git a/crates/core/src/functions.rs b/crates/core/src/functions.rs index a56873c58..69d10559b 100644 --- a/crates/core/src/functions.rs +++ b/crates/core/src/functions.rs @@ -607,7 +607,12 @@ expr_fn_vec!(named_struct); expr_fn!(from_unixtime, unixtime); expr_fn!(arrow_typeof, arg_1); expr_fn!(arrow_cast, arg_1 datatype); +expr_fn!(arrow_try_cast, arg_1 datatype); +expr_fn!(arrow_field, arg_1); +expr_fn!(cast_to_type, arg_1 reference); +expr_fn!(try_cast_to_type, arg_1 reference); expr_fn_vec!(arrow_metadata); +expr_fn_vec!(with_metadata); expr_fn!(union_tag, arg1); expr_fn!(random); @@ -966,7 +971,12 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(array_agg))?; m.add_wrapped(wrap_pyfunction!(arrow_typeof))?; m.add_wrapped(wrap_pyfunction!(arrow_cast))?; + m.add_wrapped(wrap_pyfunction!(arrow_try_cast))?; + m.add_wrapped(wrap_pyfunction!(arrow_field))?; + m.add_wrapped(wrap_pyfunction!(cast_to_type))?; + m.add_wrapped(wrap_pyfunction!(try_cast_to_type))?; m.add_wrapped(wrap_pyfunction!(arrow_metadata))?; + m.add_wrapped(wrap_pyfunction!(with_metadata))?; m.add_wrapped(wrap_pyfunction!(ascii))?; m.add_wrapped(wrap_pyfunction!(asin))?; m.add_wrapped(wrap_pyfunction!(asinh))?; diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 4fdbdc5d4..be3a14123 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -894,6 +894,28 @@ def cast(self, to: pa.DataType[Any] | type) -> Expr: return Expr(self.expr.cast(to)) + def try_cast(self, to: pa.DataType[Any] | type) -> Expr: + """Cast to a new data type, returning NULL on failure. + + Like :py:meth:`cast` but produces NULL instead of erroring when the + cast cannot be performed for a given row. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["oops"]}) + >>> result = df.select(col("a").try_cast(pa.float64()).alias("c")) + >>> result.collect_column("c")[0].as_py() is None + True + """ + if not isinstance(to, pa.DataType): + try: + to = self._to_pyarrow_types[to] + except KeyError as err: + error_msg = "Expected instance of pyarrow.DataType or builtins.type" + raise TypeError(error_msg) from err + + return Expr(self.expr.try_cast(to)) + def between(self, low: Any, high: Any, negated: bool = False) -> Expr: """Returns ``True`` if this expression is between a given range. diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 4e08fa1d9..f4eef763a 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -133,7 +133,9 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "arrays_overlap", "arrays_zip", "arrow_cast", + "arrow_field", "arrow_metadata", + "arrow_try_cast", "arrow_typeof", "ascii", "asin", @@ -151,6 +153,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "btrim", "cardinality", "case", + "cast_to_type", "cbrt", "ceil", "char_length", @@ -375,6 +378,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "translate", "trim", "trunc", + "try_cast_to_type", "union_extract", "union_tag", "upper", @@ -386,6 +390,7 @@ def _warn_expr_for_literal_arg(function_name: str, arg_name: str) -> None: "var_sample", "version", "when", + "with_metadata", ] @@ -2960,6 +2965,110 @@ def arrow_cast(expr: Expr, data_type: Expr | str | pa.DataType) -> Expr: return Expr(f.arrow_cast(expr.expr, data_type.expr)) +def arrow_try_cast(expr: Expr, data_type: Expr | str | pa.DataType) -> Expr: + """Casts an expression to a specified data type, returning NULL on failure. + + Like :py:func:`arrow_cast` but produces NULL instead of erroring when the + cast cannot be performed. The ``data_type`` may be a string in DataFusion + type syntax (for example ``"Float64"``), a ``pyarrow.DataType``, or an + ``Expr`` of string type. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["oops"]}) + >>> result = df.select( + ... dfn.functions.arrow_try_cast(dfn.col("a"), "Float64").alias("c") + ... ) + >>> result.collect_column("c")[0].as_py() is None + True + + >>> result = df.select( + ... dfn.functions.arrow_try_cast( + ... dfn.col("a"), data_type=pa.float64() + ... ).alias("c") + ... ) + >>> result.collect_column("c")[0].as_py() is None + True + """ + if isinstance(data_type, pa.DataType): + return expr.try_cast(data_type) + if isinstance(data_type, str): + data_type = Expr.string_literal(data_type) + return Expr(f.arrow_try_cast(expr.expr, data_type.expr)) + + +def arrow_field(expr: Expr) -> Expr: + """Returns the Arrow field information of an expression as a struct. + + The returned struct contains the field's name, data type, nullability, + and metadata. + + Examples: + >>> field = pa.field("val", pa.int64(), metadata={"k": "v"}) + >>> schema = pa.schema([field]) + >>> batch = pa.RecordBatch.from_arrays([pa.array([1])], schema=schema) + >>> ctx = dfn.SessionContext() + >>> df = ctx.create_dataframe([[batch]]) + >>> result = df.select( + ... dfn.functions.arrow_field(dfn.col("val")).alias("f") + ... ) + >>> out = result.collect_column("f")[0].as_py() + >>> out["name"], out["data_type"], out["nullable"], out["metadata"] + ('val', 'Int64', True, [('k', 'v')]) + """ + return Expr(f.arrow_field(expr.expr)) + + +def cast_to_type(value: Expr, type_ref: Expr) -> Expr: + """Casts ``value`` to the data type of ``type_ref``. + + Only the *type* of ``type_ref`` is used; its value is ignored. This is + useful when the target type comes from another column or expression + rather than being known up-front. Casts that fail produce an error; use + :py:func:`try_cast_to_type` for the NULL-on-failure variant. + + If the target type is known statically, prefer :py:func:`arrow_cast` + (or :py:func:`arrow_try_cast` for the NULL-on-failure variant) and + pass a type string or ``pyarrow.DataType`` directly. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1], "b": [1.0]}) + >>> result = df.select( + ... dfn.functions.cast_to_type( + ... dfn.col("a"), dfn.col("b") + ... ).alias("c") + ... ) + >>> result.collect_column("c")[0].as_py() + 1.0 + """ + return Expr(f.cast_to_type(value.expr, type_ref.expr)) + + +def try_cast_to_type(value: Expr, type_ref: Expr) -> Expr: + """Casts ``value`` to the data type of ``type_ref``, NULL on failure. + + Like :py:func:`cast_to_type`, but casts that fail produce NULL instead + of erroring. Only the *type* of ``type_ref`` is used; its value is + ignored. + + If the target type is known statically, prefer :py:func:`arrow_try_cast` + and pass a type string or ``pyarrow.DataType`` directly. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["oops"], "b": [1.0]}) + >>> result = df.select( + ... dfn.functions.try_cast_to_type( + ... dfn.col("a"), dfn.col("b") + ... ).alias("c") + ... ) + >>> result.collect_column("c")[0].as_py() is None + True + """ + return Expr(f.try_cast_to_type(value.expr, type_ref.expr)) + + def arrow_metadata(expr: Expr, key: Expr | str | None = None) -> Expr: """Returns the metadata of the input expression. @@ -2993,6 +3102,41 @@ def arrow_metadata(expr: Expr, key: Expr | str | None = None) -> Expr: return Expr(f.arrow_metadata(expr.expr, key.expr)) +def with_metadata(expr: Expr, metadata: dict[str, str]) -> Expr: + """Attaches Arrow field metadata (key/value pairs) to the input expression. + + This is the inverse of :py:func:`arrow_metadata`. Existing metadata on the + input field is preserved; new keys overwrite on collision. Keys must be + non-empty strings; empty values are allowed. + + An empty ``metadata`` dict is a no-op and returns the input expression + unchanged. Empty keys raise :py:class:`ValueError`. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.with_metadata( + ... dfn.col("a"), {"unit": "ms"} + ... ).alias("a") + ... ) + >>> result.select( + ... dfn.functions.arrow_metadata(dfn.col("a"), "unit").alias("u") + ... ).collect_column("u")[0].as_py() + 'ms' + """ + if not metadata: + return expr + args = [expr.expr] + for k, v in metadata.items(): + if not k: + msg = "with_metadata keys must be non-empty strings" + raise ValueError(msg) + args.append(Expr.string_literal(k).expr) + args.append(Expr.string_literal(v).expr) + return Expr(f.with_metadata(*args)) + + def get_field(expr: Expr, *names: Expr | str) -> Expr: """Extracts a (possibly nested) field from a struct or map by name. diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py index b02cdf6db..73a7ad727 100644 --- a/python/tests/test_functions.py +++ b/python/tests/test_functions.py @@ -1333,30 +1333,90 @@ def test_make_time(df): assert result.column(0)[0].as_py() == time(12, 30) -def test_arrow_cast(df): - df = df.select( - f.arrow_cast(column("b"), "Float64").alias("b_as_float"), - f.arrow_cast(column("b"), "Int32").alias("b_as_int"), +@pytest.mark.parametrize("cast_fn", [f.arrow_cast, f.arrow_try_cast]) +@pytest.mark.parametrize( + ("data_type", "expected"), + [ + ("Float64", pa.array([4.0, 5.0, 6.0], type=pa.float64())), + ("Int32", pa.array([4, 5, 6], type=pa.int32())), + (pa.float64(), pa.array([4.0, 5.0, 6.0], type=pa.float64())), + (pa.int32(), pa.array([4, 5, 6], type=pa.int32())), + (pa.string(), pa.array(["4", "5", "6"], type=pa.string())), + ], +) +def test_arrow_cast_variants(df, cast_fn, data_type, expected): + """arrow_cast / arrow_try_cast accept str and pyarrow target types.""" + result = df.select(cast_fn(column("b"), data_type).alias("c")).collect()[0] + assert result.column(0) == expected + + +def test_arrow_try_cast_null_on_failure(): + ctx = SessionContext() + batch = pa.RecordBatch.from_arrays([pa.array(["1.5", "oops", "3"])], names=["s"]) + df = ctx.create_dataframe([[batch]]) + + result = df.select(f.arrow_try_cast(column("s"), "Float64").alias("c")).collect()[0] + + assert result.column(0).to_pylist() == [1.5, None, 3.0] + + +def test_arrow_field(): + ctx = SessionContext() + field = pa.field("val", pa.int64(), metadata={"k": "v"}) + schema = pa.schema([field]) + batch = pa.RecordBatch.from_arrays([pa.array([1])], schema=schema) + df = ctx.create_dataframe([[batch]]) + + out = ( + df.select(f.arrow_field(column("val")).alias("f")) + .collect_column("f")[0] + .as_py() ) - result = df.collect() - assert len(result) == 1 - result = result[0] + assert out == { + "name": "val", + "data_type": "Int64", + "nullable": True, + "metadata": [("k", "v")], + } + + +@pytest.mark.parametrize( + ("cast_fn", "values", "expected"), + [ + (f.cast_to_type, pa.array([4, 5, 6]), [4.0, 5.0, 6.0]), + (f.try_cast_to_type, pa.array(["oops", "2", "3"]), [None, 2.0, 3.0]), + ], +) +def test_cast_to_type(cast_fn, values, expected): + """cast_to_type / try_cast_to_type take target type from ``type_ref``.""" + ctx = SessionContext() + batch = pa.RecordBatch.from_arrays( + [values, pa.array([1.0, 2.0, 3.0])], names=["v", "fl"] + ) + df = ctx.create_dataframe([[batch]]) - assert result.column(0) == pa.array([4.0, 5.0, 6.0], type=pa.float64()) - assert result.column(1) == pa.array([4, 5, 6], type=pa.int32()) + result = df.select(cast_fn(column("v"), column("fl")).alias("c")).collect()[0] + assert result.column(0).to_pylist() == expected + assert result.column(0).type == pa.float64() -def test_arrow_cast_with_pyarrow_type(df): - df = df.select( - f.arrow_cast(column("b"), pa.float64()).alias("b_as_float"), - f.arrow_cast(column("b"), pa.int32()).alias("b_as_int"), - f.arrow_cast(column("b"), pa.string()).alias("b_as_str"), + +def test_with_metadata_round_trip(df): + df = df.select(f.with_metadata(column("b"), {"unit": "ms"}).alias("b")) + result = df.select(f.arrow_metadata(column("b"), "unit").alias("u")).collect_column( + "u" ) - result = df.collect()[0] + assert result[0].as_py() == "ms" + + +def test_with_metadata_empty_dict_noop(df): + out = df.select(f.with_metadata(column("b"), {}).alias("b")).collect()[0] + assert out.column(0) == pa.array([4, 5, 6]) + - assert result.column(0) == pa.array([4.0, 5.0, 6.0], type=pa.float64()) - assert result.column(1) == pa.array([4, 5, 6], type=pa.int32()) - assert result.column(2) == pa.array(["4", "5", "6"], type=pa.string()) +def test_with_metadata_empty_key_raises(): + with pytest.raises(ValueError, match="non-empty"): + f.with_metadata(column("b"), {"": "v"}) def test_case(df): diff --git a/skills/datafusion_python/SKILL.md b/skills/datafusion_python/SKILL.md index 1aeb78777..4ddf54e99 100644 --- a/skills/datafusion_python/SKILL.md +++ b/skills/datafusion_python/SKILL.md @@ -758,7 +758,12 @@ F.left(col("c_phone"), lit(2)) # prefix shortcut **Hash**: `md5`, `sha224`, `sha256`, `sha384`, `sha512`, `digest` -**Type**: `arrow_typeof`, `arrow_cast`, `arrow_metadata` +**Type**: `arrow_typeof`, `arrow_cast`, `arrow_try_cast`, `arrow_field`, +`arrow_metadata`, `cast_to_type`, `with_metadata` + +Note: ``cast_to_type(value, type_ref, *, try_cast=False)`` is the single +Python entry point for both upstream ``cast_to_type`` and ``try_cast_to_type``; +pass ``try_cast=True`` for the variant that returns NULL on failure. **Other**: `in_list`, `order_by`, `alias`, `col`, `encode`, `decode`, `to_hex`, `to_char`, `uuid`, `version`, `bit_length`, `octet_length` From 55fd2c925c95711ad7222a5a48b6691d482d6056 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 16 Jun 2026 18:13:07 -0400 Subject: [PATCH 62/83] docs: convert reStructuredText sources to MyST markdown (#1579) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: convert restructuredText sources to MyST markdown Phase 2 of the documentation-site refresh. Run `rst2myst convert` over every human-authored .rst file under docs/source/ and remove the originals. The result: - 33 .rst files become 33 .md files (user guide, contributor guide, index, links). - Headings, paragraphs, hyperlinks, code blocks, admonitions, and toctree directives all map cleanly to MyST syntax. - Cross-reference anchors round-trip through MyST as `(label)=` blocks. The converter kebab-cased the labels (e.g. `(io-csv)=`), but every `{ref}` target in the corpus still uses the underscore form from the original RST (`{ref}\`CSV \``) and so do the Python docstrings that AutoAPI pulls in. Rewrite the anchors back to the underscore form so the existing references resolve. - 86 `{eval-rst}` blocks remain — they all wrap `.. ipython::` directives, which have no first-class MyST equivalent. They render identically and don't block the build. conf.py changes: - Enable `colon_fence` and `deflist` MyST extensions (rst-to-myst emits these on a few files, particularly execution-metrics.md). - Keep `.rst` in `source_suffix` even though no human-authored RST remains: sphinx-autoapi generates RST under autoapi/ at build time and Sphinx needs the suffix registered to parse it. AGENTS.md: update the two .rst paths called out under "Aggregate and Window Function Documentation" to point at the .md equivalents. Verified by building locally — `build succeeded`, no warnings, all internal cross-references resolve, the ipython examples on the landing page and basics page still execute. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: fix Apache license header format in converted markdown files RST-to-MD conversion emitted MyST `%` comment syntax with blank line between each header line, which renders as visible text. Replace with canonical `` HTML comment block matching upstream apache/datafusion and this repo's existing markdown files. Co-Authored-By: Claude Opus 4.7 (1M context) * docs: fix broken cross-reference links in distributing-work The RST -> MyST conversion left two intra-page links as undefined reference-style links, which CommonMark renders as literal bracketed text (no Sphinx warning, so the --fail-on-warning build still passed). Point both at the auto-generated heading anchors instead. Co-Authored-By: Claude Opus 4.8 (1M context) * docs: execute examples via myst-nb; native tables and validated refs Removes the last RST-syntax islands from the converted MyST markdown so the docs are markdown-native for both human and LLM authors. Executable examples (A): replace IPython.sphinxext.ipython_directive with myst-nb. The 83 `{eval-rst}` + `.. ipython:: python` blocks become native `{code-cell} ipython3` blocks, and the 14 pages that carry them gain jupytext/kernelspec front matter so myst-nb runs them. conf.py routes .md through myst-nb with nb_execution_mode="force" and nb_execution_raise_on_error=True, so a failing example now fails the build. myst-nb gives each page its own kernel instead of the IPython directive's single namespace shared across all documents in build order. That isolation surfaced expressions.md, which only ever worked by inheriting `col`/`lit` from an earlier-built page — it now imports them itself. It also changes the execution working directory to each page's own folder, so build.sh symlinks the example data next to every page that reads it by relative name and registers the python3 kernel; CI now calls build.sh so it matches local. Tables (B): the 3 `.. list-table::` directives become GFM markdown tables. Cross-references (C): the two intra-page links in distributing-work.md that the conversion left as undefined markdown references (and that built green while rendering literal brackets) become `{ref}` roles backed by explicit `(label)=` targets, so a future break fails the build instead of shipping silently. Co-Authored-By: Claude Opus 4.8 (1M context) * docs: render DataFrame cell outputs as text, not the HTML widget myst-nb prefers a cell's `_repr_html_` over its text repr. A datafusion DataFrame's HTML repr is a Jupyter-oriented widget — inline styles plus an injected