diff --git a/.ai/skills/audit-skill-md/SKILL.md b/.ai/skills/audit-skill-md/SKILL.md new file mode 100644 index 000000000..ba5255a59 --- /dev/null +++ b/.ai/skills/audit-skill-md/SKILL.md @@ -0,0 +1,290 @@ + + +--- +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/__init__.py` | "Available Functions (Categorized)", scattered uses throughout | +| `functions.spark` | `python/datafusion/functions/spark.py` | "Available Functions (Categorized)" → "Spark-Compatible Functions" subsection | +| 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/__init__.py` `__all__` and the "Available Functions (Categorized)" section | +| `spark-functions` | `functions/spark.py` `__all__`, the "Spark-Compatible Functions" subsection, and the divergent-semantics table | +| `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/__init__.py`'s `__all__`. +- Spark function names mentioned in the "Spark-Compatible Functions" + subsection should appear in `python/datafusion/functions/spark.py`'s + `__all__`. Also confirm the divergent-semantics table still matches the + current spark vs. main signatures. +- 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/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md new file mode 100644 index 000000000..a3d82a670 --- /dev/null +++ b/.ai/skills/check-upstream/SKILL.md @@ -0,0 +1,477 @@ + + +--- +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. + +**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. + +### 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()` + +**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. 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 + +**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` + +**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. 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 + +**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 +- `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 +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 + +### 8. Spark-Compatible Functions (`datafusion-spark` crate) + +**Upstream source of truth:** +- Crate source: https://github.com/apache/datafusion/tree/main/datafusion/spark/src +- Rust docs: https://docs.rs/datafusion-spark/latest/datafusion_spark/ + +**Where they are exposed in this project:** +- Python API: `python/datafusion/functions/spark.py` — each function wraps + a call to `datafusion._internal.functions.spark`; the public surface is + the module's `__all__` list. +- Rust bindings: `crates/core/src/spark_functions.rs` — `#[pyfunction]` + definitions registered via `init_module()` and re-exported under + `datafusion._internal.functions.spark`. + +**Coverage policy:** The spark namespace mirrors +`pyspark.sql.functions` parameter names and shapes exactly so pyspark +callers can paste code unchanged. Extras over pyspark are permitted as +long as positional pyspark calls still work — for example, the spark +`avg` / `try_sum` / `collect_list` / `collect_set` retain the +`distinct`/`filter`/`order_by`/`null_treatment` kwargs from the main +namespace while pyspark's single-positional form continues to work. + +**How to check:** +1. Fetch the upstream `datafusion-spark` function list from the crate + source under `datafusion/spark/src/function/` (each subdirectory is a + category: `string/`, `math/`, `datetime/`, etc.). The crate's + `function.rs` collects all `ScalarUDF` factories. +2. Cross-reference against `pyspark.sql.functions` for the public-facing + shape — pyspark is the contract this namespace is matching. +3. Compare against the functions listed in + `python/datafusion/functions/spark.py`'s `__all__`. A function is + covered if it exists in the Python `spark` namespace, even if it + aliases another function's Rust binding. +4. Report functions that are missing from the Python spark namespace. + +### 9. `__all__` Hygiene (functions.py and functions/spark.py) + +Independent of upstream parity, also flag public `def` symbols in +`python/datafusion/functions.py` **and** `python/datafusion/functions/spark.py` +that are missing from that file'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 each file 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. + +- 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/.ai/skills/make-pythonic/SKILL.md b/.ai/skills/make-pythonic/SKILL.md new file mode 100644 index 000000000..7d490ec03 --- /dev/null +++ b/.ai/skills/make-pythonic/SKILL.md @@ -0,0 +1,465 @@ + + +--- +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. + +## Scope: `functions` vs `functions.spark` + +Both `python/datafusion/functions/__init__.py` and +`python/datafusion/functions/spark.py` are in scope. We want both to feel +pythonic — accept native Python types where the argument is contextually +a literal — but `functions.spark` carries an additional constraint: +**every signature must remain compatible with `pyspark.sql.functions`**. + +Compatibility rules for the spark namespace: + +- **Parameter names must match pyspark exactly.** Pyspark callers pass by + keyword (`spark.shiftleft(col=..., numBits=...)`), so renames break + them. Do NOT rename a parameter just because it would be more pythonic + in the main namespace. +- **Positional order must match pyspark exactly.** Reordering breaks + positional pyspark calls. +- **Type unions may widen the input set, never narrow it.** Pyspark + accepts `Column` or `str` (column name) for most args; we accept + `Expr` already, and widening to `Expr | int` / `Expr | str` for + literal-friendly arguments is on-brand because the int/str case is + exactly what a pyspark caller would also try. Just verify the widened + set is a superset of what pyspark accepts for that arg. +- **Extra keyword arguments are allowed** as long as they default to + `None` and pyspark's positional/keyword form still works (e.g. the + spark `avg`/`try_sum`/`collect_list`/`collect_set` retain DataFusion's + `distinct`/`filter`/`order_by`/`null_treatment` kwargs). + +Practical effect: in `functions.spark`, apply Categories A and (where +pyspark exposes the same arg as a non-`Expr`) B normally, but cross-check +each proposed signature against `pyspark.sql.functions` before landing +it. When pyspark's own type hint is `Column | str` for a "column name" +arg, prefer leaving the spark wrapper at `Expr` — Category C +("`Expr | str` meaning column name") is unusual in `functions.py` and +should remain so in `functions.spark`. + +## 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/__init__.py` **and** `python/datafusion/functions/spark.py`. When updating a spark-namespace function, apply the compatibility rules from "Scope" above on top of the standard analysis. + +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/__init__.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/__init__.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/.cargo/config.toml b/.cargo/config.toml deleted file mode 100644 index 91a099a61..000000000 --- a/.cargo/config.toml +++ /dev/null @@ -1,12 +0,0 @@ -[target.x86_64-apple-darwin] -rustflags = [ - "-C", "link-arg=-undefined", - "-C", "link-arg=dynamic_lookup", -] - -[target.aarch64-apple-darwin] -rustflags = [ - "-C", "link-arg=-undefined", - "-C", "link-arg=dynamic_lookup", -] - 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/.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 2dc8b96fe..c35801b11 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -15,34 +15,67 @@ # specific language governing permissions and limitations # under the License. -name: Python Release Build +# Reusable workflow for running building +# This ensures the same tests run for both debug (PRs) and release (main/tags) builds + +name: Build + on: - pull_request: - branches: ["main"] - push: - tags: ["*-rc*"] - branches: ["branch-*"] + workflow_call: + inputs: + build_mode: + description: 'Build mode: debug or release' + required: true + type: string + run_wheels: + description: 'Whether to build distribution wheels' + required: false + type: boolean + default: false + +env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + UV_LOCKED: true jobs: - build: + # ============================================ + # Linting Jobs + # ============================================ + lint-rust: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Setup Rust + uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 + with: + toolchain: "nightly" + components: rustfmt + + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + + - name: Check formatting + run: cargo +nightly fmt --all -- --check + + lint-python: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 + - name: Install Python - uses: actions/setup-python@v6 + uses: actions/setup-python@v5 with: python-version: "3.12" - - uses: astral-sh/setup-uv@v7 + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 with: - enable-cache: true + enable-cache: true - # Use the --no-install-package to only install the dependencies - # but do not yet build the rust library - name: Install dependencies run: uv sync --dev --no-install-package datafusion - # Update output format to enable automatic inline annotations. - name: Run Ruff run: | uv run --no-project ruff check --output-format=github python/ @@ -50,62 +83,258 @@ jobs: - name: Run codespell run: | - uv run --no-project codespell --toml pyproject.toml + uv run --no-project codespell --toml pyproject.toml + + lint-toml: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Install taplo + uses: taiki-e/install-action@v2 + with: + tool: taplo-cli + + # if you encounter an error, try running 'taplo format' to fix the formatting automatically. + - name: Check Cargo.toml formatting + run: taplo format --check + + check-crates-patch: + if: inputs.build_mode == 'release' && startsWith(github.ref, 'refs/tags/') + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Ensure [patch.crates-io] is empty + run: python3 dev/check_crates_patch.py generate-license: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v5 - - uses: astral-sh/setup-uv@v7 + - uses: actions/checkout@v6 + + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 + with: + enable-cache: true + + - name: Install cargo-license + uses: taiki-e/install-action@v2 with: - enable-cache: true + tool: cargo-license - name: Generate license file run: uv run --no-project python ./dev/create_license.py - - uses: actions/upload-artifact@v4 + + - uses: actions/upload-artifact@v7 with: name: python-wheel-license path: LICENSE.txt + # ============================================ + # Build - Linux x86_64 + # ============================================ + build-manylinux-x86_64: + needs: [generate-license, lint-rust, lint-python] + name: Linux x86_64 (${{ matrix.python-tag }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-tag: ["abi3", "3.14t"] + steps: + - uses: actions/checkout@v6 + + - run: rm LICENSE.txt + - name: Download LICENSE.txt + uses: actions/download-artifact@v8 + with: + name: python-wheel-license + path: . + + - name: Setup Rust + uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 + + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} + + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 + 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 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" + + # 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 + manylinux: "2_28" + working-directory: examples/datafusion-ffi-example + args: --out dist + rustup-components: rust-std + + - name: Archive wheels + uses: actions/upload-artifact@v7 + with: + 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 + path: examples/datafusion-ffi-example/dist/* + + # ============================================ + # Build - Linux ARM64 + # ============================================ + build-manylinux-aarch64: + needs: [generate-license, lint-rust, lint-python] + name: Linux arm64 (${{ matrix.python-tag }}) + runs-on: ubuntu-24.04-arm + strategy: + fail-fast: false + matrix: + python-tag: ["abi3", "3.14t"] + steps: + - uses: actions/checkout@v6 + + - run: rm LICENSE.txt + - name: Download LICENSE.txt + uses: actions/download-artifact@v8 + with: + name: python-wheel-license + path: . + + - name: Setup Rust + uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 + + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + key: ${{ inputs.build_mode }}-${{ matrix.python-tag }} + + - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 + 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 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" + + - name: Archive wheels + uses: actions/upload-artifact@v7 + if: inputs.build_mode == 'release' + with: + name: dist-manylinux-aarch64-${{ matrix.python-tag }} + path: dist/* + + # ============================================ + # Build - macOS arm64 / Windows + # ============================================ build-python-mac-win: - needs: [generate-license] - name: Mac/Win + needs: [generate-license, lint-rust, lint-python] + 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.14t"] steps: - - uses: actions/checkout@v5 - - - uses: actions/setup-python@v6 - with: - python-version: ${{ matrix.python-version }} + - uses: actions/checkout@v6 - - uses: dtolnay/rust-toolchain@stable + - name: Setup Rust + uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 - run: rm LICENSE.txt - name: Download LICENSE.txt - uses: actions/download-artifact@v5 + uses: actions/download-artifact@v8 with: name: python-wheel-license path: . + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + 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@fac544c07dec837d0ccb6301d7b5580bf5edae39 + with: + enable-cache: true + - name: Install Protoc uses: arduino/setup-protoc@v3 with: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@v7 - with: - enable-cache: true + - name: Install dependencies + if: matrix.python-tag == 'abi3' + run: uv sync --dev --no-install-package datafusion - - name: Build Python package - run: | - uv sync --dev --no-install-package datafusion - uv run --no-project maturin build --release --strip --features substrait + # 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' && matrix.python-tag == 'abi3' + run: cargo clippy --no-deps --all-targets --features substrait -- -D warnings + + - 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' @@ -119,127 +348,94 @@ jobs: run: find target/wheels/ - name: Archive wheels - uses: actions/upload-artifact@v4 + 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/* + # ============================================ + # Build - macOS x86_64 (release only) + # ============================================ build-macos-x86_64: - needs: [generate-license] - name: Mac 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.14t"] steps: - - uses: actions/checkout@v5 - - - uses: actions/setup-python@v6 - with: - python-version: ${{ matrix.python-version }} + - uses: actions/checkout@v6 - - uses: dtolnay/rust-toolchain@stable + - name: Setup Rust + uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 - run: rm LICENSE.txt - name: Download LICENSE.txt - uses: actions/download-artifact@v5 + uses: actions/download-artifact@v8 with: name: python-wheel-license path: . + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + 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@fac544c07dec837d0ccb6301d7b5580bf5edae39 + with: + enable-cache: true + - name: Install Protoc uses: arduino/setup-protoc@v3 with: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@v7 - with: - enable-cache: true + - name: Install dependencies + if: matrix.python-tag == 'abi3' + run: uv sync --dev --no-install-package datafusion - - name: Build Python package - run: | - uv sync --dev --no-install-package datafusion - 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/ - name: Archive wheels - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 with: - name: dist-macos-aarch64 + name: dist-macos-aarch64-${{ matrix.python-tag }} path: target/wheels/* - build-manylinux-x86_64: - needs: [generate-license] - name: Manylinux x86_64 - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v5 - - run: rm LICENSE.txt - - name: Download LICENSE.txt - uses: actions/download-artifact@v5 - with: - name: python-wheel-license - path: . - - run: cat LICENSE.txt - - name: Build wheels - uses: PyO3/maturin-action@v1 - env: - RUST_BACKTRACE: 1 - with: - rust-toolchain: nightly - target: x86_64 - manylinux: auto - rustup-components: rust-std rustfmt # Keep them in one line due to https://github.com/PyO3/maturin-action/issues/153 - args: --release --manylinux 2014 --features protoc,substrait - - name: Archive wheels - uses: actions/upload-artifact@v4 - with: - name: dist-manylinux-x86_64 - path: target/wheels/* - - build-manylinux-aarch64: - needs: [generate-license] - name: Manylinux arm64 - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v5 - - run: rm LICENSE.txt - - name: Download LICENSE.txt - uses: actions/download-artifact@v5 - with: - name: python-wheel-license - path: . - - run: cat LICENSE.txt - - name: Build wheels - uses: PyO3/maturin-action@v1 - env: - RUST_BACKTRACE: 1 - with: - rust-toolchain: nightly - target: aarch64 - # Use manylinux_2_28-cross because the manylinux2014-cross has GCC 4.8.5, which causes the build to fail - manylinux: 2_28 - rustup-components: rust-std rustfmt # Keep them in one line due to https://github.com/PyO3/maturin-action/issues/153 - args: --release --features protoc,substrait - - name: Archive wheels - uses: actions/upload-artifact@v4 - with: - name: dist-manylinux-aarch64 - path: target/wheels/* + # ============================================ + # Build - Source Distribution + # ============================================ build-sdist: needs: [generate-license] name: Source distribution + if: inputs.build_mode == 'release' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - run: rm LICENSE.txt - name: Download LICENSE.txt - uses: actions/download-artifact@v5 + uses: actions/download-artifact@v8 with: name: python-wheel-license path: . @@ -253,16 +449,22 @@ jobs: args: --release --sdist --out dist --features protoc,substrait - name: Assert sdist build does not generate wheels run: | - if [ "$(ls -A target/wheels)" ]; then - echo "Error: Sdist build generated wheels" - exit 1 - else - echo "Directory is clean" - fi + if [ "$(ls -A target/wheels)" ]; then + echo "Error: Sdist build generated wheels" + exit 1 + else + echo "Directory is clean" + fi shell: bash - + + # ============================================ + # Build - Source Distribution + # ============================================ + merge-build-artifacts: runs-on: ubuntu-latest + name: Merge build artifacts + if: inputs.build_mode == 'release' needs: - build-python-mac-win - build-macos-x86_64 @@ -271,11 +473,14 @@ jobs: - build-sdist steps: - name: Merge Build Artifacts - uses: actions/upload-artifact/merge@v4 + uses: actions/upload-artifact/merge@v7 with: name: dist pattern: dist-* + # ============================================ + # Build - Documentation + # ============================================ # Documentation build job that runs after wheels are built build-docs: name: Build docs @@ -299,11 +504,11 @@ jobs: fi - name: Checkout docs sources - uses: actions/checkout@v5 + uses: actions/checkout@v6 - name: Checkout docs target branch if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') - uses: actions/checkout@v5 + uses: actions/checkout@v6 with: fetch-depth: 0 ref: ${{ steps.target-branch.outputs.value }} @@ -312,34 +517,36 @@ jobs: - name: Setup Python uses: actions/setup-python@v6 with: - python-version: "3.11" + python-version: "3.10" - name: Install dependencies - uses: astral-sh/setup-uv@v7 + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 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@v5 + uses: actions/download-artifact@v8 with: - name: dist-manylinux-x86_64 + name: dist-manylinux-x86_64-abi3 path: wheels/ - # Install from the pre-built wheel - - name: Install from pre-built wheel + # Install from the pre-built wheels + - name: Install from pre-built wheels run: | set -x uv venv # Install documentation dependencies uv sync --dev --no-install-package datafusion --group docs - # Install the pre-built wheel - WHEEL=$(find wheels/ -name "*.whl" | head -1) - if [ -n "$WHEEL" ]; then - echo "Installing wheel: $WHEEL" - uv pip install "$WHEEL" + # Install all pre-built wheels + WHEELS=$(find wheels/ -name "*.whl") + if [ -n "$WHEELS" ]; then + echo "Installing wheels:" + echo "$WHEELS" + uv pip install wheels/*.whl else - echo "ERROR: No wheel found!" + echo "ERROR: No wheels found!" exit 1 fi @@ -347,9 +554,10 @@ jobs: run: | set -x cd docs - curl -O https://gist.githubusercontent.com/ritchie46/cac6b337ea52281aa23c049250a4ff03/raw/89a957ff3919d90e6ef2d34235e6bf22304f3366/pokemon.csv - curl -O https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2021-01.parquet - uv run --no-project make html + # build.sh downloads the example data, registers the Jupyter kernel + # myst-nb needs, symlinks the data next to each executed page, and + # runs sphinx. Using it here keeps CI identical to a local build. + uv run --no-project bash ./build.sh - name: Copy & push the generated HTML if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') @@ -368,16 +576,3 @@ jobs: git commit -m 'Publish built docs triggered by ${{ github.sha }}' git push || git push --force fi - - # NOTE: PyPI publish needs to be done manually for now after release passed the vote - # release: - # name: Publish in PyPI - # needs: [build-manylinux, build-python-mac-win] - # runs-on: ubuntu-latest - # steps: - # - uses: actions/download-artifact@v5 - # - name: Publish to PyPI - # uses: pypa/gh-action-pypi-publish@master - # with: - # user: __token__ - # password: ${{ secrets.pypi_password }} diff --git a/python/datafusion/html_formatter.py b/.github/workflows/ci.yml similarity index 61% rename from python/datafusion/html_formatter.py rename to .github/workflows/ci.yml index 65eb1f042..ab284b522 100644 --- a/python/datafusion/html_formatter.py +++ b/.github/workflows/ci.yml @@ -15,15 +15,27 @@ # specific language governing permissions and limitations # under the License. -"""Deprecated module for dataframe formatting.""" +# CI workflow for pull requests - runs tests in DEBUG mode for faster feedback -import warnings +name: CI -from datafusion.dataframe_formatter import * # noqa: F403 +on: + pull_request: + branches: ["main"] -warnings.warn( - "The module 'html_formatter' is deprecated and will be removed in the next release." - "Please use 'dataframe_formatter' instead.", - DeprecationWarning, - stacklevel=3, -) +concurrency: + group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + build: + uses: ./.github/workflows/build.yml + with: + build_mode: debug + run_wheels: false + secrets: inherit + + test: + needs: build + uses: ./.github/workflows/test.yml + secrets: inherit diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml new file mode 100644 index 000000000..2d0f166ba --- /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@8aad20d150bbac5944a9f9d289da16a4b0d87c1e # v4 + with: + languages: actions + + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@8aad20d150bbac5944a9f9d289da16a4b0d87c1e # v4 + with: + category: "/language:actions" diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index ac45e9fdf..841bf205a 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -16,7 +16,12 @@ # under the License. name: Dev -on: [push, pull_request] +on: + push: + branches: + - main + - branch-* + pull_request: jobs: @@ -25,10 +30,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v5 + uses: actions/checkout@v6 - name: Setup Python uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.14" - name: Audit licenses run: ./dev/release/run-rat.sh . diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 000000000..bddc89eac --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,49 @@ +# 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. + +# Release workflow - runs tests in RELEASE mode and builds distribution wheels +# Triggered on: +# - Merges to main +# - Release candidate tags (*-rc*) +# - Release tags (e.g., 45.0.0) + +name: Release Build + +on: + push: + branches: + - "main" + tags: + - "*-rc*" # Release candidates (e.g., 45.0.0-rc1) + - "[0-9]+.*" # Release tags (e.g., 45.0.0) + +concurrency: + group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + build: + uses: ./.github/workflows/build.yml + with: + build_mode: release + run_wheels: true + secrets: inherit + + test: + needs: build + uses: ./.github/workflows/test.yml + secrets: inherit diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml deleted file mode 100644 index 99457b872..000000000 --- a/.github/workflows/test.yaml +++ /dev/null @@ -1,139 +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. - -name: Python test -on: - push: - branches: [main] - pull_request: - branches: [main] - -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -jobs: - test-matrix: - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - python-version: - - "3.10" - - "3.11" - - "3.12" - - "3.13" - - "3.14" - toolchain: - - "stable" - - steps: - - uses: actions/checkout@v5 - - - name: Verify example datafusion version - run: | - MAIN_VERSION=$(grep -A 1 "name = \"datafusion-common\"" Cargo.lock | grep "version = " | head -1 | sed 's/.*version = "\(.*\)"/\1/') - EXAMPLE_VERSION=$(grep -A 1 "name = \"datafusion-common\"" examples/datafusion-ffi-example/Cargo.lock | grep "version = " | head -1 | sed 's/.*version = "\(.*\)"/\1/') - echo "Main crate datafusion version: $MAIN_VERSION" - echo "FFI example datafusion version: $EXAMPLE_VERSION" - - if [ "$MAIN_VERSION" != "$EXAMPLE_VERSION" ]; then - echo "❌ Error: FFI example datafusion versions don't match!" - exit 1 - fi - - - name: Setup Rust Toolchain - uses: dtolnay/rust-toolchain@stable - id: rust-toolchain - with: - components: clippy,rustfmt - - - name: Install Protoc - uses: arduino/setup-protoc@v3 - with: - version: '27.4' - repo-token: ${{ secrets.GITHUB_TOKEN }} - - - name: Setup Python - uses: actions/setup-python@v6 - with: - python-version: ${{ matrix.python-version }} - - - name: Cache Cargo - uses: actions/cache@v4 - with: - path: ~/.cargo - key: cargo-cache-${{ steps.rust-toolchain.outputs.cachekey }}-${{ hashFiles('Cargo.lock') }} - - - name: Run Clippy - if: ${{ matrix.python-version == '3.10' && matrix.toolchain == 'stable' }} - run: cargo clippy --all-targets --all-features -- -D clippy::all -D warnings -A clippy::redundant_closure - - - name: Install dependencies and build - uses: astral-sh/setup-uv@v7 - with: - enable-cache: true - - - name: Run tests - env: - RUST_BACKTRACE: 1 - run: | - git submodule update --init - uv sync --dev --no-install-package datafusion - uv run --no-project maturin develop --uv - uv run --no-project pytest -v . - - - name: FFI unit tests - run: | - cd examples/datafusion-ffi-example - uv run --no-project maturin develop --uv - uv run --no-project pytest python/tests/_test*.py - - - name: Cache the generated dataset - id: cache-tpch-dataset - uses: actions/cache@v4 - with: - path: benchmarks/tpch/data - key: tpch-data-2.18.0 - - - name: Run dbgen to create 1 Gb dataset - if: ${{ steps.cache-tpch-dataset.outputs.cache-hit != 'true' }} - run: | - cd benchmarks/tpch - RUN_IN_CI=TRUE ./tpch-gen.sh 1 - - - name: Run TPC-H examples - run: | - cd examples/tpch - uv run --no-project python convert_data_to_parquet.py - uv run --no-project pytest _tests.py - - nightly-fmt: - runs-on: ubuntu-latest - - steps: - - uses: actions/checkout@v5 - - - name: Setup Rust Toolchain - uses: dtolnay/rust-toolchain@stable - id: rust-toolchain - with: - toolchain: "nightly" - components: clippy,rustfmt - - - name: Check Formatting - run: cargo +nightly fmt -- --check diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 000000000..558e751c8 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,137 @@ +# 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. + +# Reusable workflow for running tests. +# Single matrix covers both GIL (abi3 wheel) and free-threaded +# (per-interpreter wheels) builds. + +name: Test + +on: + workflow_call: + +env: + UV_LOCKED: true + +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: + 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.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-stable-${{ hashFiles('Cargo.lock') }} + + - name: Install dependencies + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 + with: + enable-cache: true + + - name: Download pre-built Linux wheel + uses: actions/download-artifact@v8 + with: + name: dist-manylinux-x86_64-${{ matrix.wheel-tag }} + path: wheels/ + + # 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/ + + - name: Install from pre-built wheels + run: | + set -x + # 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.14t environment as the 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 --python "$VENV_PY" wheels/*.whl + else + echo "ERROR: No wheels found!" + exit 1 + fi + + - 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 + # 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 + uv pip install tpchgen-cli + 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/.github/workflows/verify-release-candidate.yml b/.github/workflows/verify-release-candidate.yml new file mode 100644 index 000000000..6ecb547b5 --- /dev/null +++ b/.github/workflows/verify-release-candidate.yml @@ -0,0 +1,83 @@ +# 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: Verify Release Candidate + +# NOTE: This workflow is intended to be run manually via workflow_dispatch. + +on: + workflow_dispatch: + inputs: + version: + description: Version number (e.g., 52.0.0) + required: true + type: string + rc_number: + description: Release candidate number (e.g., 1) + required: true + type: string + +concurrency: + group: ${{ github.repository }}-${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + verify: + name: Verify RC (${{ matrix.os }}-${{ matrix.arch }}) + strategy: + fail-fast: false + matrix: + include: + # Linux + - os: linux + arch: x64 + runner: ubuntu-latest + - os: linux + arch: arm64 + runner: ubuntu-24.04-arm + + # macOS + - os: macos + arch: arm64 + runner: macos-latest + - os: macos + arch: x64 + runner: macos-15-intel + + # Windows + - os: windows + arch: x64 + runner: windows-latest + runs-on: ${{ matrix.runner }} + steps: + - name: Checkout repository + uses: actions/checkout@v6 + + - name: Set up protoc + uses: arduino/setup-protoc@v3 + with: + 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 }}" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bcefa405d..0a212480b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,12 +17,12 @@ 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 # Ruff version. - rev: v0.9.10 + rev: v0.15.1 hooks: # Run the linter. - id: ruff @@ -53,5 +53,12 @@ repos: additional_dependencies: - tomli + - repo: https://github.com/astral-sh/uv-pre-commit + # uv version. + rev: 0.10.7 + hooks: + # Update the uv lockfile + - id: uv-lock + default_language_version: python: python3 diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..fda08b23c --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,92 @@ + + +# 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`](skills/datafusion_python/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. + +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. + +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 +`.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. + +- **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. + +## 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.md` — + add new aggregate functions to the "Aggregate Functions" list and include usage + examples if appropriate. +- **Window functions**: `docs/source/user-guide/common-operations/windows.md` — + add new window functions to the "Available Functions" list and include usage + examples if appropriate. 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/Cargo.lock b/Cargo.lock index ac691ca9d..ab222b177 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,54 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 - -[[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", -] +version = 4 [[package]] name = "adler2" @@ -87,9 +39,9 @@ checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" [[package]] name = "alloc-stdlib" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" dependencies = [ "alloc-no-stdlib", ] @@ -111,53 +63,24 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.100" +version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" - -[[package]] -name = "apache-avro" -version = "0.20.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a033b4ced7c585199fb78ef50fca7fe2f444369ec48080c5fd072efa1a03cc7" -dependencies = [ - "bigdecimal", - "bon", - "bzip2 0.6.1", - "crc32fast", - "digest", - "log", - "miniz_oxide", - "num-bigint", - "quad-rand", - "rand", - "regex-lite", - "serde", - "serde_bytes", - "serde_json", - "snap", - "strum", - "strum_macros", - "thiserror", - "uuid", - "xz2", - "zstd", -] +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" [[package]] name = "ar_archive_writer" -version = "0.2.0" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0c269894b6fe5e9d7ada0cf69b5bf847ff35bc25fc271f08e1d080fce80339a" +checksum = "4087686b4b0a3427190bae57a1d9a478dbb2d40c5dc1bd6e2b6d797913bdd348" dependencies = [ "object", ] [[package]] name = "arc-swap" -version = "1.8.0" +version = "1.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51d03449bb8ca2cc2ef70869af31463d1ae5ccc8fa3e334b307203fbf815207e" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" dependencies = [ "rustversion", ] @@ -176,9 +99,9 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb372a7cbcac02a35d3fb7b3fc1f969ec078e871f9bb899bf00a2e1809bec8a3" +checksum = "61d285d16bce7d0be61912f7928342b673067b6b7d7ef6cc179258ba7de1fecf" dependencies = [ "arrow-arith", "arrow-array", @@ -198,9 +121,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f377dcd19e440174596d83deb49cd724886d91060c07fec4f67014ef9d54049" +checksum = "757ef1836251e88222542a7da2623bc1c9cb9e20afefa6db2c41e79991cd91d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -212,9 +135,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a23eaff85a44e9fa914660fb0d0bb00b79c4a3d888b5334adb3ea4330c84f002" +checksum = "bc9a4a4b2b5ecd0e04df03471661cb61f28bed3c7fd50994715129b01b2edb97" dependencies = [ "ahash", "arrow-buffer", @@ -223,29 +146,54 @@ dependencies = [ "chrono", "chrono-tz", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", + "libc", "num-complex", "num-integer", "num-traits", ] +[[package]] +name = "arrow-avro" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb45cd6bd2b25c0965793b83200eaca82214273a8030fbbc2d783e4c7c65a61" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-schema", + "bytes", + "bzip2", + "crc", + "flate2", + "indexmap", + "liblzma", + "rand 0.9.4", + "serde", + "serde_json", + "snap", + "strum_macros", + "uuid", + "zstd", +] + [[package]] name = "arrow-buffer" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2819d893750cb3380ab31ebdc8c68874dd4429f90fd09180f3c93538bd21626" +checksum = "c12b576ef18c1deb80925a248b25ad84f419198d791b8e293fc6aaa60441fe90" dependencies = [ "bytes", "half", - "num-bigint", + "num-bigint 0.5.1", "num-traits", ] [[package]] name = "arrow-cast" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3d131abb183f80c450d4591dc784f8d7750c50c6e2bc3fcaad148afc8361271" +checksum = "68338a9096a5dc9bc11927c58c43a8526d96bf6abd2012ef6c0c9f505991cc79" dependencies = [ "arrow-array", "arrow-buffer", @@ -254,7 +202,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64", + "base64 0.23.1", "chrono", "comfy-table", "half", @@ -265,9 +213,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2275877a0e5e7e7c76954669366c2aa1a829e340ab1f612e647507860906fb6b" +checksum = "25011b52b346407d497ef0030e12b45e4f2d0cc279efc09c4f3d09106db30e36" dependencies = [ "arrow-array", "arrow-cast", @@ -280,9 +228,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05738f3d42cb922b9096f7786f606fcb8669260c2640df8490533bb2fa38c9d3" +checksum = "723fe4aeed7604e00b9883a465af4ff0a0e6c44c03e41a68c3d1cbc403e0e44d" dependencies = [ "arrow-buffer", "arrow-schema", @@ -293,9 +241,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d09446e8076c4b3f235603d9ea7c5494e73d441b01cd61fb33d7254c11964b3" +checksum = "149437b14371f5b9ec60f5ddc751483ae99d7a7072653c0075e5e469156eea7b" dependencies = [ "arrow-array", "arrow-buffer", @@ -309,15 +257,16 @@ dependencies = [ [[package]] name = "arrow-json" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "371ffd66fa77f71d7628c63f209c9ca5341081051aa32f9c8020feb0def787c0" +checksum = "f18b9123ccfec418a663f821c9a034af339711678c11ffe00d3ec07da5ff9f7e" dependencies = [ "arrow-array", "arrow-buffer", "arrow-cast", - "arrow-data", + "arrow-ord", "arrow-schema", + "arrow-select", "chrono", "half", "indexmap", @@ -333,9 +282,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc94fc7adec5d1ba9e8cd1b1e8d6f72423b33fe978bf1f46d970fafab787521" +checksum = "e6c08dff0686cf23ca4f562803f191ccbeb726dbae6309cd4b4aaf65e0f2c979" dependencies = [ "arrow-array", "arrow-buffer", @@ -346,9 +295,9 @@ dependencies = [ [[package]] name = "arrow-pyarrow" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbd810e3997bae72f58cda57231ccb0a2fda07911ca1b0a5718cbf9379abb297" +checksum = "c196ecc25b3a8dcbc1d842f2619cee653dcfa2fb8b56a291bc0481c3cf5c3821" dependencies = [ "arrow-array", "arrow-data", @@ -358,9 +307,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "169676f317157dc079cc5def6354d16db63d8861d61046d2f3883268ced6f99f" +checksum = "bbec439386df71ad570e6758a946111322b9e9dc8db83b5527321f0b4c9119c2" dependencies = [ "arrow-array", "arrow-buffer", @@ -371,9 +320,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d27609cd7dd45f006abae27995c2729ef6f4b9361cde1ddd019dc31a5aa017e0" +checksum = "e6fed2ca0d1eade57e811cbe73b98ad50cc08a1183e13b2d2aa43a7df593f40e" dependencies = [ "bitflags", "serde_core", @@ -382,9 +331,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae980d021879ea119dd6e2a13912d81e64abed372d53163e804dfe84639d8010" +checksum = "466b19cf75130b891dc1b23a84b343c714c62c64c9c62e365c76aa0ff90a53fb" dependencies = [ "ahash", "arrow-array", @@ -396,9 +345,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf35e8ef49dcf0c5f6d175edee6b8af7b45611805333129c541a8b89a0fc0534" +checksum = "c838a25bb3691e919e0f617616ac51a4ff8517a952e29ca133cf0c22b2ce65b1" dependencies = [ "arrow-array", "arrow-buffer", @@ -411,33 +360,16 @@ 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.19" +version = "0.4.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06575e6a9673580f52661c92107baabffbf41e2141373441cbcdc47cb733003c" +checksum = "e79b3f8a79cccc2898f31920fc69f304859b3bd567490f75ebf51ae1c792a9ac" dependencies = [ - "bzip2 0.5.2", - "flate2", - "futures-core", - "memchr", + "compression-codecs", + "compression-core", "pin-project-lite", "tokio", - "xz2", - "zstd", - "zstd-safe", ] [[package]] @@ -445,9 +377,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" @@ -457,7 +386,7 @@ checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -468,7 +397,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -488,9 +417,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" @@ -498,6 +427,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "bigdecimal" version = "0.4.10" @@ -506,17 +441,16 @@ checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695" dependencies = [ "autocfg", "libm", - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", - "serde", ] [[package]] name = "bitflags" -version = "2.10.0" +version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "812e12b5285cc515a9c72a5c1d3b6d46a19dac5acfef5265968c166106e31dd3" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" [[package]] name = "blake2" @@ -524,20 +458,21 @@ version = "0.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" dependencies = [ - "digest", + "digest 0.10.7", ] [[package]] name = "blake3" -version = "1.8.2" +version = "1.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3888aaa89e4b2a40fca9848e400f6a658a5a3978de7be858e209cafa8be9a4a0" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" dependencies = [ "arrayref", "arrayvec", "cc", "cfg-if", "constant_time_eq", + "cpufeatures", ] [[package]] @@ -550,35 +485,19 @@ dependencies = [ ] [[package]] -name = "bon" -version = "3.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebeb9aaf9329dff6ceb65c689ca3db33dbf15f324909c60e4e5eef5701ce31b1" -dependencies = [ - "bon-macros", - "rustversion", -] - -[[package]] -name = "bon-macros" -version = "3.8.1" +name = "block-buffer" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77e9d642a7e3a318e37c2c9427b5a6a48aa1ad55dcd986f3034ab2239045a645" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" dependencies = [ - "darling", - "ident_case", - "prettyplease", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.113", + "hybrid-array", ] [[package]] name = "brotli" -version = "8.0.2" +version = "8.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -587,9 +506,9 @@ dependencies = [ [[package]] name = "brotli-decompressor" -version = "5.0.0" +version = "5.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -597,30 +516,15 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.19.1" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" - -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "bytes" -version = "1.11.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" - -[[package]] -name = "bzip2" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47" -dependencies = [ - "bzip2-sys", -] +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" [[package]] name = "bzip2" @@ -631,21 +535,11 @@ dependencies = [ "libbz2-rs-sys", ] -[[package]] -name = "bzip2-sys" -version = "0.1.13+1.0.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" -dependencies = [ - "cc", - "pkg-config", -] - [[package]] name = "cc" -version = "1.2.51" +version = "1.2.64" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a0aeaff4ff1a90589618835a598e545176939b97874f7abc7851caa0618f203" +checksum = "dad887fd958be91b5098c0248def011f4523ab786cd411be668777e55063501f" dependencies = [ "find-msvc-tools", "jobserver", @@ -665,11 +559,22 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core 0.10.1", +] + [[package]] name = "chrono" -version = "0.4.42" +version = "0.4.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "145052bdd345b87320e369255277e3fb5152762ad123a901ef5c262dd38fe8d2" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" dependencies = [ "iana-time-zone", "num-traits", @@ -689,23 +594,50 @@ dependencies = [ [[package]] name = "cmake" -version = "0.1.57" +version = "0.1.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" dependencies = [ "cc", ] [[package]] name = "comfy-table" -version = "7.2.1" +version = "7.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b03b7db8e0b4b2fdad6c551e634134e99ec000e5c8c3b6856c65e8bbaded7a3b" +checksum = "958c5d6ecf1f214b4c2bbbbf6ab9523a864bd136dcf71a7e8904799acfe1ad47" dependencies = [ "unicode-segmentation", "unicode-width", ] +[[package]] +name = "compression-codecs" +version = "0.4.38" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce2548391e9c1929c21bf6aa2680af86fe4c1b33e6cea9ac1cfeec0bd11218cf" +dependencies = [ + "bzip2", + "compression-core", + "flate2", + "liblzma", + "memchr", + "zstd", + "zstd-safe", +] + +[[package]] +name = "compression-core" +version = "0.4.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc14f565cf027a105f7a44ccf9e5b424348421a1d8952a8fc9d499d313107789" + +[[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" @@ -721,25 +653,16 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" dependencies = [ - "getrandom 0.2.16", + "getrandom 0.2.17", "once_cell", "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.3.1" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" [[package]] name = "core-foundation" @@ -758,28 +681,28 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" [[package]] -name = "core_extensions" -version = "1.5.4" +name = "cpufeatures" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42bb5e5d0269fd4f739ea6cedaf29c16d81c27a7ce7582008e90eb50dcd57003" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" dependencies = [ - "core_extensions_proc_macros", + "libc", ] [[package]] -name = "core_extensions_proc_macros" -version = "1.5.4" +name = "crc" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "533d38ecd2709b7608fb8e18e4504deb99e9a72879e6aa66373a76d8dc4259ea" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] [[package]] -name = "cpufeatures" -version = "0.2.17" +name = "crc-catalog" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" -dependencies = [ - "libc", -] +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" [[package]] name = "crc32fast" @@ -790,15 +713,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "crossbeam-channel" -version = "0.5.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" -dependencies = [ - "crossbeam-utils", -] - [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -821,6 +735,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + [[package]] name = "cstr" version = "0.2.12" @@ -852,46 +775,11 @@ dependencies = [ "memchr", ] -[[package]] -name = "darling" -version = "0.21.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" -dependencies = [ - "darling_core", - "darling_macro", -] - -[[package]] -name = "darling_core" -version = "0.21.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" -dependencies = [ - "fnv", - "ident_case", - "proc-macro2", - "quote", - "strsim", - "syn 2.0.113", -] - -[[package]] -name = "darling_macro" -version = "0.21.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" -dependencies = [ - "darling_core", - "quote", - "syn 2.0.113", -] - [[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", @@ -903,15 +791,13 @@ dependencies = [ [[package]] name = "datafusion" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ba7cb113e9c0bedf9e9765926031e132fa05a1b09ba6e93a6d1a4d7044457b8" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "arrow-schema", "async-trait", - "bytes", - "bzip2 0.6.1", + "bzip2", "chrono", "datafusion-catalog", "datafusion-catalog-listing", @@ -941,28 +827,25 @@ dependencies = [ "datafusion-sql", "flate2", "futures", - "itertools", + "indexmap", + "itertools 0.15.0", + "liblzma", "log", "object_store", "parking_lot", "parquet", - "rand", - "regex", - "rstest", "sqlparser", "tempfile", "tokio", "url", "uuid", - "xz2", "zstd", ] [[package]] name = "datafusion-catalog" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66a3a799f914a59b1ea343906a0486f17061f39509af74e874a866428951130d" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "async-trait", @@ -976,7 +859,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools", + "itertools 0.15.0", "log", "object_store", "parking_lot", @@ -985,9 +868,8 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db1b113c80d7a0febcd901476a57aef378e717c54517a163ed51417d87621b0" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "async-trait", @@ -1001,42 +883,42 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "futures", - "itertools", + "itertools 0.15.0", "log", "object_store", - "tokio", + "percent-encoding", ] [[package]] name = "datafusion-common" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c10f7659e96127d25e8366be7c8be4109595d6a2c3eac70421f380a7006a1b0" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ - "ahash", - "apache-avro", "arrow", "arrow-ipc", + "arrow-schema", "chrono", + "foldhash 0.2.0", "half", - "hashbrown 0.14.5", + "hashbrown 0.17.1", "indexmap", + "itertools 0.15.0", "libc", "log", + "num-traits", "object_store", "parquet", - "paste", "recursive", "sqlparser", "tokio", + "uuid", "web-time", ] [[package]] name = "datafusion-common-runtime" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b92065bbc6532c6651e2f7dd30b55cba0c7a14f860c7e1d15f165c41a1868d95" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "futures", "log", @@ -1045,15 +927,14 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fde13794244bc7581cd82f6fff217068ed79cdc344cafe4ab2c3a1c3510b38d6" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "async-compression", "async-trait", "bytes", - "bzip2 0.6.1", + "bzip2", "chrono", "datafusion-common", "datafusion-common-runtime", @@ -1063,26 +944,27 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "flate2", "futures", "glob", - "itertools", + "itertools 0.15.0", + "liblzma", "log", "object_store", - "rand", + "parking_lot", + "rand 0.9.4", "tokio", "tokio-util", "url", - "xz2", "zstd", ] [[package]] name = "datafusion-datasource-arrow" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "804fa9b4ecf3157982021770617200ef7c1b2979d57bec9044748314775a9aea" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "arrow-ipc", @@ -1095,38 +977,36 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", - "itertools", + "itertools 0.15.0", "object_store", "tokio", ] [[package]] name = "datafusion-datasource-avro" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "388ed8be535f562cc655b9c3d22edbfb0f1a50a25c242647a98b6d92a75b55a1" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" 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 = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61a1641a40b259bab38131c5e6f48fac0717bedb7dc93690e604142a849e0568" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "async-trait", @@ -1138,6 +1018,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1147,9 +1028,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adeacdb00c1d37271176f8fb6a1d8ce096baba16ea7a4b2671840c5c9c64fe85" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "async-trait", @@ -1161,19 +1041,21 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", "tokio", + "tokio-stream", ] [[package]] name = "datafusion-datasource-parquet" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43d0b60ffd66f28bfb026565d62b0a6cbc416da09814766a3797bba7d85a3cd9" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", + "arrow-schema", "async-trait", "bytes", "datafusion-common", @@ -1181,15 +1063,17 @@ dependencies = [ "datafusion-datasource", "datafusion-execution", "datafusion-expr", + "datafusion-functions", "datafusion-functions-aggregate-common", "datafusion-physical-expr", "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-pruning", "datafusion-session", "futures", - "itertools", + "itertools 0.15.0", "log", "object_store", "parking_lot", @@ -1199,37 +1083,41 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b99e13947667b36ad713549237362afb054b2d8f8cc447751e23ec61202db07" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" [[package]] name = "datafusion-execution" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63695643190679037bc946ad46a263b62016931547bf119859c511f7ff2f5178" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", + "arrow-buffer", "async-trait", + "bytes", "dashmap", "datafusion-common", "datafusion-expr", + "datafusion-physical-expr-common", "futures", "log", "object_store", "parking_lot", - "rand", + "pin-project-lite", + "rand 0.9.4", "tempfile", + "tokio", + "tokio-util", "url", ] [[package]] name = "datafusion-expr" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9a4787cbf5feb1ab351f789063398f67654a6df75c4d37d7f637dc96f951a91" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", + "arrow-schema", "async-trait", "chrono", "datafusion-common", @@ -1238,9 +1126,10 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-functions-window-common", "datafusion-physical-expr-common", + "datafusion-proto-common", + "datafusion-proto-models", "indexmap", - "itertools", - "paste", + "itertools 0.15.0", "recursive", "serde_json", "sqlparser", @@ -1248,77 +1137,105 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ce2fb1b8c15c9ac45b0863c30b268c69dc9ee7a1ee13ecf5d067738338173dc" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "datafusion-common", "indexmap", - "itertools", - "paste", + "itertools 0.15.0", ] [[package]] name = "datafusion-ffi" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec510e7787641279b0336e8b79e4b7bd1385d5976875ff9b97f4269ce5231a67" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ - "abi_stable", "arrow", "arrow-schema", "async-ffi", "async-trait", - "datafusion", + "chrono", + "datafusion-catalog", "datafusion-common", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", "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", ] +[[package]] +name = "datafusion-ffi-example" +version = "54.0.0" +dependencies = [ + "arrow", + "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", + "pyo3-log", +] + [[package]] name = "datafusion-functions" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "794a9db7f7b96b3346fc007ff25e994f09b8f0511b4cf7dff651fadfe3ebb28f" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "arrow-buffer", - "base64", + "base64 0.23.1", "blake2", "blake3", "chrono", + "chrono-tz", "datafusion-common", "datafusion-doc", "datafusion-execution", "datafusion-expr", "datafusion-expr-common", "datafusion-macros", + "datafusion-physical-expr-common", "hex", - "itertools", + "itertools 0.15.0", "log", - "md-5", + "md-5 0.11.0", + "memchr", "num-traits", - "rand", + "rand 0.9.4", "regex", "sha2", - "unicode-segmentation", "uuid", ] [[package]] name = "datafusion-functions-aggregate" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c25210520a9dcf9c2b2cbbce31ebd4131ef5af7fc60ee92b266dc7d159cb305" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-doc", @@ -1329,17 +1246,16 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "half", + "hashbrown 0.17.1", "log", - "paste", + "num-traits", ] [[package]] name = "datafusion-functions-aggregate-common" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62f4a66f3b87300bb70f4124b55434d2ae3fe80455f3574701d0348da040b55d" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr-common", @@ -1348,9 +1264,8 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae5c06eed03918dc7fe7a9f082a284050f0e9ecf95d72f57712d1496da03b8c4" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "arrow-ord", @@ -1364,32 +1279,32 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-macros", "datafusion-physical-expr-common", - "itertools", + "hashbrown 0.17.1", + "itertools 0.15.0", + "itoa", "log", - "paste", + "memchr", ] [[package]] name = "datafusion-functions-table" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db4fed1d71738fbe22e2712d71396db04c25de4111f1ec252b8f4c6d3b25d7f5" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "async-trait", "datafusion-catalog", "datafusion-common", "datafusion-expr", + "datafusion-physical-expr", "datafusion-physical-plan", "parking_lot", - "paste", ] [[package]] name = "datafusion-functions-window" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d92206aa5ae21892f1552b4d61758a862a70956e6fd7a95cb85db1de74bc6d1" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "datafusion-common", @@ -1400,14 +1315,12 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "log", - "paste", ] [[package]] name = "datafusion-functions-window-common" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53ae9bcc39800820d53a22d758b3b8726ff84a5a3e24cecef04ef4e5fdf1c7cc" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1415,20 +1328,18 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1063ad4c9e094b3f798acee16d9a47bd7372d9699be2de21b05c3bd3f34ab848" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "datafusion-doc", "quote", - "syn 2.0.113", + "syn 3.0.3", ] [[package]] name = "datafusion-optimizer" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f35f9ec5d08b87fd1893a30c2929f2559c2f9806ca072d8fefca5009dc0f06a" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "chrono", @@ -1437,7 +1348,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-physical-expr", "indexmap", - "itertools", + "itertools 0.15.0", "log", "recursive", "regex", @@ -1446,31 +1357,30 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c30cc8012e9eedcb48bbe112c6eff4ae5ed19cf3003cb0f505662e88b7014c5d" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr", "datafusion-expr-common", "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", + "datafusion-proto-models", "half", - "hashbrown 0.14.5", + "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.15.0", "parking_lot", - "paste", - "petgraph 0.8.3", + "petgraph", + "recursive", + "tokio", ] [[package]] name = "datafusion-physical-expr-adapter" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f9ff2dbd476221b1f67337699eff432781c4e6e1713d2aefdaa517dfbf79768" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "datafusion-common", @@ -1478,28 +1388,30 @@ dependencies = [ "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-common", - "itertools", + "itertools 0.15.0", ] [[package]] name = "datafusion-physical-expr-common" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90da43e1ec550b172f34c87ec68161986ced70fd05c8d2a2add66eef9c276f03" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ - "ahash", "arrow", + "chrono", "datafusion-common", "datafusion-expr-common", - "hashbrown 0.14.5", - "itertools", + "datafusion-proto-models", + "hashbrown 0.17.1", + "indexmap", + "itertools 0.15.0", + "parking_lot", + "pin-project", ] [[package]] name = "datafusion-physical-optimizer" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce9804f799acd7daef3be7aaffe77c0033768ed8fdbf5fb82fc4c5f2e6bc14e6" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "datafusion-common", @@ -1510,49 +1422,53 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-pruning", - "itertools", + "datafusion-session", + "itertools 0.15.0", "recursive", ] [[package]] name = "datafusion-physical-plan" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0acf0ad6b6924c6b1aa7d213b181e012e2d3ec0a64ff5b10ee6282ab0f8532ac" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ - "ahash", "arrow", + "arrow-data", + "arrow-ipc", "arrow-ord", "arrow-schema", "async-trait", - "chrono", + "bytes", "datafusion-common", "datafusion-common-runtime", "datafusion-execution", "datafusion-expr", + "datafusion-functions", "datafusion-functions-aggregate-common", "datafusion-functions-window-common", "datafusion-physical-expr", "datafusion-physical-expr-common", + "datafusion-proto-common", + "datafusion-proto-models", "futures", "half", - "hashbrown 0.14.5", + "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.15.0", "log", + "num-traits", "parking_lot", "pin-project-lite", + "serde_json", "tokio", ] [[package]] name = "datafusion-proto" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d368093a98a17d1449b1083ac22ed16b7128e4c67789991869480d8c4a40ecb9" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", - "chrono", "datafusion-catalog", "datafusion-catalog-listing", "datafusion-common", @@ -1568,26 +1484,35 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-proto-common", + "datafusion-proto-models", "object_store", "prost", ] [[package]] name = "datafusion-proto-common" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b6aef3d5e5c1d2bc3114c4876730cb76a9bdc5a8df31ef1b6db48f0c1671895" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "datafusion-common", "prost", ] +[[package]] +name = "datafusion-proto-models" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +dependencies = [ + "datafusion-common", + "datafusion-proto-common", + "prost", +] + [[package]] name = "datafusion-pruning" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac2c2498a1f134a9e11a9f5ed202a2a7d7e9774bd9249295593053ea3be999db" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "datafusion-common", @@ -1596,21 +1521,23 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", - "itertools", "log", ] [[package]] name = "datafusion-python" -version = "51.0.0" +version = "54.0.0" dependencies = [ "arrow", "arrow-select", "async-trait", + "chrono", "cstr", "datafusion", "datafusion-ffi", "datafusion-proto", + "datafusion-python-util", + "datafusion-spark", "datafusion-substrait", "futures", "log", @@ -1623,17 +1550,31 @@ dependencies = [ "pyo3-async-runtimes", "pyo3-build-config", "pyo3-log", + "serde_json", "tokio", "url", "uuid", ] +[[package]] +name = "datafusion-python-util" +version = "54.0.0" +dependencies = [ + "arrow", + "datafusion", + "datafusion-ffi", + "datafusion-proto", + "prost", + "pyo3", + "tokio", +] + [[package]] name = "datafusion-session" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f96eebd17555386f459037c65ab73aae8df09f464524c709d6a3134ad4f4776" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ + "arrow-schema", "async-trait", "datafusion-common", "datafusion-execution", @@ -1642,43 +1583,71 @@ dependencies = [ "parking_lot", ] +[[package]] +name = "datafusion-spark" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +dependencies = [ + "arrow", + "bigdecimal", + "chrono", + "crc32fast", + "datafusion", + "datafusion-catalog", + "datafusion-common", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions", + "datafusion-functions-aggregate", + "datafusion-functions-aggregate-common", + "datafusion-functions-nested", + "log", + "num-traits", + "percent-encoding", + "rand 0.9.4", + "serde_json", + "sha1", + "sha2", + "twox-hash", + "url", +] + [[package]] name = "datafusion-sql" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fc195fe60634b2c6ccfd131b487de46dc30eccae8a3c35a13f136e7f440414f" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "arrow", "bigdecimal", "chrono", "datafusion-common", "datafusion-expr", + "datafusion-functions-nested", "indexmap", "log", "recursive", "regex", "sqlparser", + "stacker", ] [[package]] name = "datafusion-substrait" -version = "51.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2505af06d103a55b4e8ded0c6aeb6c72a771948da939c0bd3f8eee67af475a9c" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" dependencies = [ "async-recursion", "async-trait", "chrono", "datafusion", "half", - "itertools", + "itertools 0.15.0", "object_store", "pbjson-types", "prost", "substrait", "tokio", "url", - "uuid", ] [[package]] @@ -1687,20 +1656,31 @@ 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.1", + "const-oid", + "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", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -1711,9 +1691,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" @@ -1733,15 +1713,15 @@ 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" -version = "0.1.6" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "645cbb3a84e60b7531617d5ae4e57f7e27308f6445f5abf653209ea76dec8dff" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" [[package]] name = "fixedbitset" @@ -1761,13 +1741,13 @@ dependencies = [ [[package]] name = "flate2" -version = "1.1.5" +version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfe33edd8e85a12a67454e37f8c75e730830d83e313556ab9ebf9ee7fbeb3bfb" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", - "libz-rs-sys", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -1799,9 +1779,9 @@ dependencies = [ [[package]] name = "futures" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" dependencies = [ "futures-channel", "futures-core", @@ -1814,9 +1794,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", "futures-sink", @@ -1824,15 +1804,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" [[package]] name = "futures-executor" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" dependencies = [ "futures-core", "futures-task", @@ -1841,44 +1821,38 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" [[package]] name = "futures-macro" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] name = "futures-sink" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" [[package]] name = "futures-task" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" - -[[package]] -name = "futures-timer" -version = "3.0.3" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" [[package]] name = "futures-util" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-channel", "futures-core", @@ -1888,19 +1862,9 @@ dependencies = [ "futures-task", "memchr", "pin-project-lite", - "pin-utils", "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" @@ -1913,9 +1877,9 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.2.16" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", "js-sys", @@ -1933,11 +1897,23 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 5.3.0", "wasip2", "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", + "rand_core 0.10.1", +] + [[package]] name = "glob" version = "0.3.3" @@ -1946,9 +1922,9 @@ checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "h2" -version = "0.4.12" +version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c0b69cfcb4e1b9f1bf2f53f95f766e4661169728ec61cd3fe5a0166f2d1386" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" dependencies = [ "atomic-waker", "bytes", @@ -1980,10 +1956,6 @@ name = "hashbrown" version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" -dependencies = [ - "ahash", - "allocator-api2", -] [[package]] name = "hashbrown" @@ -2005,6 +1977,17 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +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" version = "0.5.0" @@ -2019,9 +2002,9 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "http" -version = "1.4.0" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" dependencies = [ "bytes", "itoa", @@ -2062,11 +2045,20 @@ 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" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" dependencies = [ "atomic-waker", "bytes", @@ -2078,7 +2070,6 @@ dependencies = [ "httparse", "itoa", "pin-project-lite", - "pin-utils", "smallvec", "tokio", "want", @@ -2086,16 +2077,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", @@ -2103,14 +2093,13 @@ dependencies = [ [[package]] name = "hyper-util" -version = "0.1.19" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "727805d60e7938b76b826a6ef209eb70eaa1812794f9424d4a4e2d740662df5f" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", - "futures-core", "futures-util", "http", "http-body", @@ -2127,9 +2116,9 @@ dependencies = [ [[package]] name = "iana-time-zone" -version = "0.1.64" +version = "0.1.65" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33e57f83510bb73707521ebaffa789ec8caf86f9657cad665b092b581d40e9fb" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" dependencies = [ "android_system_properties", "core-foundation-sys", @@ -2151,12 +2140,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", @@ -2164,9 +2154,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", @@ -2177,9 +2167,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", @@ -2191,15 +2181,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", @@ -2211,15 +2201,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", @@ -2230,12 +2220,6 @@ dependencies = [ "zerovec", ] -[[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" @@ -2249,9 +2233,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", @@ -2259,59 +2243,45 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.12.1" +version = "2.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ad4bb2b565bca0645f4d68c5c9af97fba094e9791da685bf83cb5f3ce74acf2" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.16.1", -] - -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", + "hashbrown 0.17.1", + "serde", + "serde_core", ] -[[package]] -name = "integer-encoding" -version = "3.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" - [[package]] name = "ipnet" -version = "2.11.0" +version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "469fb0b9cefa57e3ef31275ee7cacb78f2fdca44e4765491884a2b119d4eb130" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" [[package]] -name = "iri-string" -version = "0.7.10" +name = "itertools" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c91338f0783edbd6195decb37bae672fd3b165faffb89bf7b9e6942f8b1a731a" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" dependencies = [ - "memchr", - "serde", + "either", ] [[package]] name = "itertools" -version = "0.14.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" dependencies = [ "either", ] [[package]] name = "itoa" -version = "1.0.17" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jobserver" @@ -2325,11 +2295,12 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.83" +version = "0.3.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8" +checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31" dependencies = [ - "once_cell", + "cfg-if", + "futures-util", "wasm-bindgen", ] @@ -2392,62 +2363,72 @@ 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" -version = "0.2.179" +version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5a2d376baa530d1238d133232d15e239abad80d05838b4b59354e5268af431f" +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]] -name = "libm" -version = "0.2.15" +name = "liblzma" +version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" +checksum = "b6033b77c21d1f56deeae8014eb9fbe7bdf1765185a6c508b5ca82eeaed7f899" +dependencies = [ + "liblzma-sys", +] [[package]] -name = "libmimalloc-sys" -version = "0.1.44" +name = "liblzma-sys" +version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" +checksum = "1a60851d15cd8c5346eca4ab8babff585be2ae4bc8097c067291d3ffe2add3b6" dependencies = [ "cc", "libc", + "pkg-config", ] [[package]] -name = "libz-rs-sys" -version = "0.5.5" +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c10501e7805cee23da17c7790e59df2870c0d4043ec6d03f67d31e2b53e77415" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" dependencies = [ - "zlib-rs", + "cc", ] [[package]] name = "linux-raw-sys" -version = "0.11.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" +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" @@ -2460,9 +2441,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.29" +version = "0.4.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" [[package]] name = "lru-slab" @@ -2472,54 +2453,44 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.12.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab6473172471198271ff72e9379150e9dfd70d8e533e0752a27e515b48dd375e" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" dependencies = [ "twox-hash", ] [[package]] -name = "lzma-sys" -version = "0.1.20" +name = "md-5" +version = "0.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" +checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" dependencies = [ - "cc", - "libc", - "pkg-config", + "cfg-if", + "digest 0.10.7", ] [[package]] name = "md-5" -version = "0.10.6" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" dependencies = [ "cfg-if", - "digest", + "digest 0.11.3", ] [[package]] name = "memchr" -version = "2.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" - -[[package]] -name = "memoffset" -version = "0.9.1" +version = "2.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" [[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", ] @@ -2536,9 +2507,9 @@ dependencies = [ [[package]] name = "mio" -version = "1.1.1" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "wasi", @@ -2559,7 +2530,16 @@ checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" dependencies = [ "num-integer", "num-traits", - "serde", +] + +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", ] [[package]] @@ -2592,39 +2572,41 @@ dependencies = [ [[package]] name = "object" -version = "0.32.2" +version = "0.37.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6a622008b6e321afc04970976f62ee297fdbaa6f95318ca343e3eebb9648441" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" dependencies = [ "memchr", ] [[package]] name = "object_store" -version = "0.12.4" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c1be0c6c22ec0817cdc77d3842f721a17fd30ab6965001415b5402a74e6b740" +checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "bytes", "chrono", "form_urlencoded", - "futures", + "futures-channel", + "futures-core", + "futures-util", "http", "http-body-util", "httparse", "humantime", "hyper", - "itertools", - "md-5", + "itertools 0.14.0", + "md-5 0.10.6", "parking_lot", "percent-encoding", "quick-xml", - "rand", + "rand 0.10.1", "reqwest", "ring", - "rustls-pemfile", + "rustls-pki-types", "serde", "serde_json", "serde_urlencoded", @@ -2638,25 +2620,16 @@ dependencies = [ ] [[package]] -name = "once_cell" -version = "1.21.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" - -[[package]] -name = "openssl-probe" -version = "0.2.0" +name = "once_cell" +version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f50d9b3dabb09ecd771ad0aa242ca6894994c130308ca3d7684634df8037391" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" [[package]] -name = "ordered-float" -version = "2.10.1" +name = "openssl-probe" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" -dependencies = [ - "num-traits", -] +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" [[package]] name = "parking_lot" @@ -2683,54 +2656,45 @@ dependencies = [ [[package]] name = "parquet" -version = "57.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be3e4f6d320dd92bfa7d612e265d7d08bba0a240bab86af3425e1d255a511d89" +checksum = "7065842956a20c2a536924ce8e4d9955f7422451511b9eb7500d7bfe5077e59c" dependencies = [ "ahash", "arrow-array", "arrow-buffer", - "arrow-cast", "arrow-data", "arrow-ipc", "arrow-schema", "arrow-select", - "base64", + "base64 0.23.1", "brotli", "bytes", "chrono", "flate2", "futures", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "lz4_flex", - "num-bigint", + "num-bigint 0.5.1", "num-integer", "num-traits", "object_store", - "paste", "seq-macro", "simdutf8", "snap", - "thrift", "tokio", "twox-hash", "zstd", ] -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - [[package]] name = "pbjson" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "898bac3fa00d0ba57a4e8289837e965baa2dee8c3749f3b11d45a64b4223d9c3" dependencies = [ - "base64", + "base64 0.22.1", "serde", ] @@ -2741,7 +2705,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af22d08a625a2213a78dbb0ffa253318c5c79ce3133d32d296655a7bdfb02095" dependencies = [ "heck", - "itertools", + "itertools 0.14.0", "prost", "prost-types", ] @@ -2767,16 +2731,6 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" -[[package]] -name = "petgraph" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" -dependencies = [ - "fixedbitset", - "indexmap", -] - [[package]] name = "petgraph" version = "0.8.3" @@ -2808,34 +2762,48 @@ dependencies = [ ] [[package]] -name = "pin-project-lite" -version = "0.2.16" +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] [[package]] -name = "pin-utils" -version = "0.1.0" +name = "pin-project-lite" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[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" -version = "1.13.0" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f89776e4d69bb58bc6993e99ffa1d11f228b839984854c7daeb5d37f87cbe950" +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", ] @@ -2856,32 +2824,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] name = "proc-macro-crate" -version = "3.4.0" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "219cb19e96be00ab2e37d6e299658a0cfa83e52429179969b0f0121b4ac46983" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" dependencies = [ "toml_edit", ] [[package]] name = "proc-macro2" -version = "1.0.104" +version = "1.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9695f8df41bb4f3d222c95a67532365f569318332d03d5f3f67f37b20e6ebdf0" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" dependencies = [ "unicode-ident", ] [[package]] name = "prost" -version = "0.14.1" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7231bd9b3d3d33c86b58adbac74b5ec0ad9f496b19d22801d773636feaa95f3d" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" dependencies = [ "bytes", "prost-derive", @@ -2889,42 +2857,41 @@ dependencies = [ [[package]] name = "prost-build" -version = "0.14.1" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" +checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ "heck", - "itertools", + "itertools 0.14.0", "log", "multimap", - "once_cell", - "petgraph 0.7.1", + "petgraph", "prettyplease", "prost", "prost-types", "regex", - "syn 2.0.113", + "syn 2.0.118", "tempfile", ] [[package]] name = "prost-derive" -version = "0.14.1" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools", + "itertools 0.14.0", "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] name = "prost-types" -version = "0.14.1" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9b4db3d6da204ed77bb26ba83b6122a73aeb2e87e25fbf7ad2e84c4ccbf8f72" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" dependencies = [ "prost", ] @@ -2940,9 +2907,9 @@ dependencies = [ [[package]] name = "psm" -version = "0.1.28" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d11f2fedc3b7dafdc2851bc52f277377c5473d378859be234bc7ebb593144d01" +checksum = "645dbe486e346d9b5de3ef16ede18c26e6c70ad97418f4874b8b1889d6e761ea" dependencies = [ "ar_archive_writer", "cc", @@ -2950,28 +2917,26 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.26.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ba0117f4212101ee6544044dae45abe1083d30ce7b29c4b5cbdfa2354e07383" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-async-runtimes" -version = "0.26.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6ee6d4cb3e8d5b925f5cdb38da183e0ff18122eb2048d4041c9e7034d026e23" +checksum = "b3ef68daa7316a3fac65e5e18b2203f010346de1c1c53456811a2624673ab046" dependencies = [ - "futures", + "futures-channel", + "futures-util", "once_cell", "pin-project-lite", "pyo3", @@ -2980,18 +2945,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.26.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fc6ddaf24947d12a9aa31ac65431fb1b851b8f4365426e182901eabfb87df5f" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.26.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "025474d3928738efb38ac36d4744a74a400c901c7596199e20e45d98eb194105" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", "pyo3-build-config", @@ -2999,9 +2964,9 @@ dependencies = [ [[package]] name = "pyo3-log" -version = "0.13.2" +version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f8bae9ad5ba08b0b0ed2bb9c2bdbaeccc69cafca96d78cf0fbcea0d45d122bb" +checksum = "f64083bd3a16a353d9d62335808e8e13d0552d2a2b83fdb084496192dcfa9fcd" dependencies = [ "arc-swap", "log", @@ -3010,40 +2975,33 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.26.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e64eb489f22fe1c95911b77c44cc41e7c19f3082fc81cce90f657cdc42ffded" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] name = "pyo3-macros-backend" -version = "0.26.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "100246c0ecf400b475341b8455a9213344569af29a3c841d29270e53102e0fcf" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", - "syn 2.0.113", + "syn 2.0.118", ] -[[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.38.4" +version = "0.39.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b66c2058c55a409d601666cffe35f04333cf1013010882cec174a7467cd4e21c" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" dependencies = [ "memchr", "serde", @@ -3071,14 +3029,14 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.13" +version = "0.11.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31" +checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" dependencies = [ "bytes", "getrandom 0.3.4", "lru-slab", - "rand", + "rand 0.9.4", "ring", "rustc-hash", "rustls", @@ -3106,9 +3064,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.42" +version = "1.0.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a338cc41d27e6cc6dce6cefc13a0729dfbb81c262b1f519331575dd80ef3067f" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" dependencies = [ "proc-macro2", ] @@ -3119,14 +3077,31 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +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", + "rand_core 0.9.5", +] + +[[package]] +name = "rand" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", ] [[package]] @@ -3136,18 +3111,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", ] [[package]] name = "rand_core" -version = "0.9.3" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "recursive" version = "0.1.1" @@ -3165,7 +3146,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -3179,9 +3160,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.12.2" +version = "1.12.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4" +checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" dependencies = [ "aho-corasick", "memchr", @@ -3191,26 +3172,20 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.13" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5276caf25ac86c8d810222b3dbb938e512c55c6831a10f3e6ed1c93b84041f1c" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" dependencies = [ "aho-corasick", "memchr", "regex-syntax", ] -[[package]] -name = "regex-lite" -version = "0.1.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d942b98df5e658f56f20d592c7f868833fe38115e65c33003d8cd224b0155da" - [[package]] name = "regex-syntax" -version = "0.8.8" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "regress" @@ -3222,28 +3197,13 @@ dependencies = [ "memchr", ] -[[package]] -name = "relative-path" -version = "1.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" - -[[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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-core", "futures-util", @@ -3287,46 +3247,17 @@ checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" dependencies = [ "cc", "cfg-if", - "getrandom 0.2.16", + "getrandom 0.2.17", "libc", "untrusted", "windows-sys 0.52.0", ] -[[package]] -name = "rstest" -version = "0.26.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" -dependencies = [ - "futures-timer", - "futures-util", - "rstest_macros", -] - -[[package]] -name = "rstest_macros" -version = "0.26.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" -dependencies = [ - "cfg-if", - "glob", - "proc-macro-crate", - "proc-macro2", - "quote", - "regex", - "relative-path", - "rustc_version", - "syn 2.0.113", - "unicode-ident", -] - [[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" @@ -3339,9 +3270,9 @@ dependencies = [ [[package]] name = "rustix" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ "bitflags", "errno", @@ -3352,9 +3283,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.35" +version = "0.23.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "533f54bc6a7d4f647e46ad909549eda97bf5afc1585190ef692b4286b198bd8f" +checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" dependencies = [ "once_cell", "ring", @@ -3366,9 +3297,9 @@ dependencies = [ [[package]] name = "rustls-native-certs" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" dependencies = [ "openssl-probe", "rustls-pki-types", @@ -3376,20 +3307,11 @@ dependencies = [ "security-framework", ] -[[package]] -name = "rustls-pemfile" -version = "2.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" -dependencies = [ - "rustls-pki-types", -] - [[package]] name = "rustls-pki-types" -version = "1.13.2" +version = "1.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21e6f2ab2928ca4291b86736a8bd920a277a399bba1589409d72154ff87c1282" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" dependencies = [ "web-time", "zeroize", @@ -3397,9 +3319,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.8" +version = "0.103.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ffdfa2f5286e2247234e03f680868ac2815974dc39e00ea15adc445d0aafe52" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" dependencies = [ "ring", "rustls-pki-types", @@ -3414,9 +3336,9 @@ checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "ryu" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a50f4cf475b65d88e057964e0e9bb1f0aa9bbb2036dc65c64596b42932536984" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "same-file" @@ -3429,9 +3351,9 @@ dependencies = [ [[package]] name = "schannel" -version = "0.1.28" +version = "0.1.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "891d81b926048e76efe18581bf793546b4c0eaf8448d72be8de2bbee5fd166e1" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" dependencies = [ "windows-sys 0.61.2", ] @@ -3457,7 +3379,7 @@ dependencies = [ "proc-macro2", "quote", "serde_derive_internals", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -3468,9 +3390,9 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "security-framework" -version = "3.5.1" +version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3297343eaf830f66ede390ea39da1d462b6b0c1b000f420d0a83f898bbbe6ef" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ "bitflags", "core-foundation", @@ -3481,9 +3403,9 @@ dependencies = [ [[package]] name = "security-framework-sys" -version = "2.15.0" +version = "2.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc1f0cbffaac4852523ce30d8bd3c5cdc873501d96ff467ca09b6767bb8cd5c0" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" dependencies = [ "core-foundation-sys", "libc", @@ -3491,9 +3413,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", @@ -3515,16 +3437,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" @@ -3542,7 +3454,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -3553,15 +3465,16 @@ checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] name = "serde_json" -version = "1.0.148" +version = "1.0.150" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3084b546a1dd6289475996f182a22aba973866ea8e8b02c51d9f46b1336a22da" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ + "indexmap", "itoa", "memchr", "serde", @@ -3571,14 +3484,14 @@ dependencies = [ [[package]] name = "serde_tokenstream" -version = "0.2.2" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64060d864397305347a78851c51588fd283767e7e7589829e8121d65512340f1" +checksum = "d7c49585c52c01f13c5c2ebb333f14f6885d76daa768d8a037d28017ec538c69" dependencies = [ "proc-macro2", "quote", "serde", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -3606,28 +3519,45 @@ dependencies = [ "unsafe-libyaml", ] +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest 0.11.3", +] + [[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", - "digest", + "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" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" [[package]] name = "simd-adler32" -version = "0.3.8" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" [[package]] name = "simdutf8" @@ -3637,21 +3567,21 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" [[package]] name = "siphasher" -version = "1.0.1" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56199f7ddabf13fe5074ce809e7d3f42b42ae711800501b5b16ea82ad029c39d" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" [[package]] name = "slab" -version = "0.4.11" +version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a2ae44ef20feb57a68b23d846850f861394c2e02dc425a50098ae8c90267589" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" [[package]] name = "smallvec" -version = "1.15.1" +version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" [[package]] name = "snap" @@ -3661,19 +3591,19 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" [[package]] name = "socket2" -version = "0.6.1" +version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17129e116933cf371d018bb80ae557e889637989d8638274fb25622827b03881" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] name = "sqlparser" -version = "0.59.0" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4591acadbcf52f0af60eafbb2c003232b2b4cd8de5f0e9437cb8b1b59046cc0f" +checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec" dependencies = [ "log", "recursive", @@ -3682,13 +3612,47 @@ dependencies = [ [[package]] name = "sqlparser_derive" -version = "0.3.0" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "stabby" +version = "72.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7b834ec7ced12095fea1e4b07dcb7e8cf2b59b18afa3eac52494d835965a5ec" +dependencies = [ + "rustversion", + "stabby-abi", +] + +[[package]] +name = "stabby-abi" +version = "72.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff1a4f477858a5bdf927c9fab7f579899de9b13e39f8b3b3b300c89fbab632f4" +dependencies = [ + "rustc_version", + "rustversion", + "sha2-const-stable", + "stabby-macros", +] + +[[package]] +name = "stabby-macros" +version = "72.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da5fc6819faabb412da764b99d3b713bb55083c11e7e0c00144d386cd6a1939c" +checksum = "b31c4b2434980b67ad83f300a58088ba14d59454dcd79ba3d87419bbd924d31e" dependencies = [ + "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -3699,48 +3663,37 @@ checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "stacker" -version = "0.1.22" +version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1f8b29fb42aafcea4edeeb6b2f2d7ecd0d969c48b4cf0d2e64aafc471dd6e59" +checksum = "640c8cdd92b6b12f5bcb1803ca3bbf5ab96e5e6b6b96b9ab77dabe9e880b3190" dependencies = [ "cc", "cfg-if", "libc", "psm", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] -[[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", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[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", @@ -3755,7 +3708,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", - "syn 2.0.113", + "syn 2.0.118", "typify", "walkdir", ] @@ -3768,9 +3721,9 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "1.0.109" +version = "2.0.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" dependencies = [ "proc-macro2", "quote", @@ -3779,9 +3732,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.113" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "678faa00651c9eb72dd2020cbdf275d92eccb2400d568e419efdd64838145cb4" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" dependencies = [ "proc-macro2", "quote", @@ -3805,23 +3758,23 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] name = "target-lexicon" -version = "0.13.4" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1dd07eb858a2067e2f3c7155d54e929265c264e6f37efe3ee7a8d1b5a1dd0ba" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "tempfile" -version = "3.24.0" +version = "3.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "655da9c7eb6305c55742045d5a8d2037996d61d8de95806335c7c86ce0f82e9c" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.2", @@ -3829,33 +3782,22 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.17" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f63587ca0f12b72a0600bcba1d40081f830876000bb46dd2337a3051618f4fc8" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.17" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ff15c8ecd7de3849db632e14d18d2571fa09dfc5ed93479bc4485c7a517c913" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", -] - -[[package]] -name = "thrift" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" -dependencies = [ - "byteorder", - "integer-encoding", - "ordered-float", + "syn 2.0.118", ] [[package]] @@ -3869,9 +3811,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", @@ -3879,9 +3821,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.10.0" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa5fdc3bce6191a1dbc8c02d5c8bffcf557bafa17c124c5264a458f1b0613fa" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" dependencies = [ "tinyvec_macros", ] @@ -3894,9 +3836,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.49.0" +version = "1.52.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "libc", @@ -3909,13 +3851,13 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.6.0" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af407857209536a95c8e56f8231ef2c2e2aff839b22e07a1ffcbc617e9db9fa5" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -3928,6 +3870,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -3943,18 +3897,18 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "0.7.5+spec-1.1.0" +version = "1.1.1+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" dependencies = [ "serde_core", ] [[package]] name = "toml_edit" -version = "0.23.10+spec-1.0.0" +version = "0.25.12+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "84c8b9f757e028cee9fa244aea147aab2a9ec09d5325a9b01e0a49730c2b5269" +checksum = "d2153edc6955a6c354fad8f5efd38b6a8769bdccf9fe50f8e1329f81b0baa5d7" dependencies = [ "indexmap", "toml_datetime", @@ -3964,18 +3918,18 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.0.6+spec-1.1.0" +version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3198b4b0a8e11f09dd03e133c0280504d0801269e9afa46362ffde1cbeebf44" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" dependencies = [ "winnow", ] [[package]] name = "tower" -version = "0.5.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", @@ -3988,20 +3942,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]] @@ -4035,7 +3989,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -4053,44 +4007,20 @@ 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" +dependencies = [ + "rand 0.9.4", +] [[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.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8c1ae7cc0fdb8b842d65d127cb981574b0d2b249b74d1c7a2986863dc134f71" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" [[package]] name = "typify" @@ -4117,7 +4047,7 @@ dependencies = [ "semver", "serde", "serde_json", - "syn 2.0.113", + "syn 2.0.118", "thiserror", "unicode-ident", ] @@ -4135,21 +4065,21 @@ dependencies = [ "serde", "serde_json", "serde_tokenstream", - "syn 2.0.113", + "syn 2.0.118", "typify-impl", ] [[package]] name = "unicode-ident" -version = "1.0.22" +version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-segmentation" -version = "1.12.0" +version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" [[package]] name = "unicode-width" @@ -4157,12 +4087,6 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "unsafe-libyaml" version = "0.2.11" @@ -4177,9 +4101,9 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" [[package]] name = "url" -version = "2.5.7" +version = "2.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08bc136a29a3d1758e07a9cca267be308aeebf5cfd5a10f3f67ab2097683ef5b" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" dependencies = [ "form_urlencoded", "idna", @@ -4195,13 +4119,12 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.19.0" +version = "1.23.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2e054861b4bd027cd373e18e8d8d8e6548085000e41290d95ce0c373a654b4a" +checksum = "144d6b123cef80b301b8f72a9e2ca4370ddec21950d0a103dd22c437006d2db7" dependencies = [ - "getrandom 0.3.4", + "getrandom 0.4.3", "js-sys", - "serde_core", "wasm-bindgen", ] @@ -4238,18 +4161,18 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] name = "wasip2" -version = "1.0.1+wasi-0.2.4" +version = "1.0.4+wasi-0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" dependencies = [ "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.106" +version = "0.2.125" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd" +checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a" dependencies = [ "cfg-if", "once_cell", @@ -4260,22 +4183,19 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.56" +version = "0.4.75" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "836d9622d604feee9e5de25ac10e3ea5f2d65b41eac0d9ce72eb5deae707ce7c" +checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280" dependencies = [ - "cfg-if", "js-sys", - "once_cell", "wasm-bindgen", - "web-sys", ] [[package]] name = "wasm-bindgen-macro" -version = "0.2.106" +version = "0.2.125" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3" +checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4283,22 +4203,22 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.106" +version = "0.2.125" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40" +checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd" dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.106" +version = "0.2.125" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4" +checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f" dependencies = [ "unicode-ident", ] @@ -4318,9 +4238,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.83" +version = "0.3.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b32828d774c412041098d182a8b38b16ea816958e07cf40eec2bc080ae137ac" +checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d" dependencies = [ "js-sys", "wasm-bindgen", @@ -4336,22 +4256,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" @@ -4361,12 +4265,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" @@ -4388,7 +4286,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -4399,7 +4297,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] @@ -4435,15 +4333,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" @@ -4593,39 +4482,30 @@ checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" [[package]] name = "winnow" -version = "0.7.14" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a5364e9d77fcdeeaa6062ced926ee3381faa2ee02d3eb83a5c27a8825540829" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" dependencies = [ "memchr", ] [[package]] name = "wit-bindgen" -version = "0.46.0" +version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" [[package]] name = "writeable" -version = "0.6.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" - -[[package]] -name = "xz2" -version = "0.1.7" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" -dependencies = [ - "lzma-sys", -] +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" [[package]] name = "yoke" -version = "0.8.1" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" dependencies = [ "stable_deref_trait", "yoke-derive", @@ -4634,68 +4514,68 @@ 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", - "syn 2.0.113", + "syn 2.0.118", "synstructure", ] [[package]] name = "zerocopy" -version = "0.8.31" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd74ec98b9250adb3ca554bdde269adf631549f51d8a8f8f0a10b50f1cb298c3" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.31" +version = "0.8.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8a8d209fdf45cf5138cbb5a506f6b52522a25afccc534d1475dad8e31105c6a" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" dependencies = [ "proc-macro2", "quote", - "syn 2.0.113", + "syn 2.0.118", ] [[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", - "syn 2.0.113", + "syn 2.0.118", "synstructure", ] [[package]] name = "zeroize" -version = "1.8.2" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" [[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", @@ -4704,9 +4584,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", @@ -4715,26 +4595,26 @@ 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", - "syn 2.0.113", + "syn 2.0.118", ] [[package]] name = "zlib-rs" -version = "0.5.5" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40990edd51aae2c2b6907af74ffb635029d5788228222c4bb811e9351c0caad3" +checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" [[package]] name = "zmij" -version = "1.0.10" +version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30e0d8dffbae3d840f64bda38e28391faef673a7b5a6017840f2a106c8145868" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" [[package]] name = "zstd" diff --git a/Cargo.toml b/Cargo.toml index 364713964..a9e15d7e9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,78 +15,70 @@ # specific language governing permissions and limitations # under the License. -[package] -name = "datafusion-python" -version = "51.0.0" +[workspace.package] +version = "54.0.0" homepage = "https://datafusion.apache.org/python" repository = "https://github.com/apache/datafusion-python" authors = ["Apache DataFusion "] description = "Apache DataFusion DataFrame and SQL Query Engine" readme = "README.md" license = "Apache-2.0" -edition = "2021" -rust-version = "1.78" -include = [ - "/src", - "/datafusion", - "/LICENSE.txt", - "build.rs", - "pyproject.toml", - "Cargo.toml", - "Cargo.lock", -] +edition = "2024" +rust-version = "1.88" -[features] -default = ["mimalloc"] -protoc = ["datafusion-substrait/protoc"] -substrait = ["dep:datafusion-substrait"] +[workspace] +members = ["crates/core", "crates/util", "examples/datafusion-ffi-example"] +resolver = "3" -[dependencies] -tokio = { version = "1.47", features = [ - "macros", - "rt", - "rt-multi-thread", - "sync", -] } -pyo3 = { version = "0.26", features = [ - "extension-module", - "abi3", - "abi3-py310", -] } -pyo3-async-runtimes = { version = "0.26", features = ["tokio-runtime"] } -pyo3-log = "0.13.2" -arrow = { version = "57", features = ["pyarrow"] } -arrow-select = { version = "57" } -datafusion = { version = "51", features = ["avro", "unicode_expressions"] } -datafusion-substrait = { version = "51", optional = true } -datafusion-proto = { version = "51" } -datafusion-ffi = { version = "51" } -prost = "0.14.1" # keep in line with `datafusion-substrait` -uuid = { version = "1.18", features = ["v4"] } -mimalloc = { version = "0.1", optional = true, default-features = false, features = [ - "local_dynamic_tls", -] } +[workspace.dependencies] +tokio = { version = "1.52" } +pyo3 = { version = "0.29" } +pyo3-async-runtimes = { version = "0.29" } +pyo3-log = "0.13.3" +chrono = { version = "0.4", default-features = false } +arrow = { version = "59" } +arrow-array = { version = "59" } +arrow-schema = { version = "59" } +arrow-select = { version = "59" } +datafusion = { version = "55.0.0" } +datafusion-substrait = { version = "55.0.0" } +datafusion-proto = { version = "55.0.0" } +datafusion-ffi = { version = "55.0.0" } +datafusion-catalog = { version = "55.0.0", default-features = false } +datafusion-common = { version = "55.0.0", default-features = false } +datafusion-functions-aggregate = { version = "55.0.0" } +datafusion-functions-window = { version = "55.0.0" } +datafusion-spark = { version = "55.0.0" } +datafusion-expr = { version = "55.0.0" } +prost = "0.14.3" +serde_json = "1" +uuid = { version = "1.23" } +mimalloc = { version = "0.1", default-features = false } async-trait = "0.1.89" futures = "0.3" cstr = "0.2" -object_store = { version = "0.12.4", features = [ - "aws", - "gcp", - "azure", - "http", -] } +object_store = { version = "0.13.1" } url = "2" -log = "0.4.27" +log = "0.4.29" parking_lot = "0.12" - -[build-dependencies] -prost-types = "0.14.1" # keep in line with `datafusion-substrait` -pyo3-build-config = "0.26" - -[lib] -name = "datafusion_python" -crate-type = ["cdylib", "rlib"] +prost-types = "0.14.3" # keep in line with `datafusion-substrait` +pyo3-build-config = "0.29" +datafusion-python-util = { path = "crates/util", version = "54.0.0" } [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. +[patch.crates-io] +datafusion = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-common = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } diff --git a/README.md b/README.md index 0cdf17ab8..f6ee662d0 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`](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 +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/skills/datafusion_python/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 @@ -275,7 +291,16 @@ needing to activate the virtual environment: ```bash uv run --no-project maturin develop --uv -uv --no-project pytest . +uv run --no-project pytest +``` + +To run the FFI tests within the examples folder, after you have built +`datafusion-python` with the previous commands: + +```bash +cd examples/datafusion-ffi-example +uv run --no-project maturin develop --uv +uv run --no-project pytest python/tests/_test_*py ``` ### Running & Installing pre-commit hooks @@ -303,6 +328,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 diff --git a/benchmarks/tpch/tpch.py b/benchmarks/tpch/tpch.py index 9cc897e76..ffee5554c 100644 --- a/benchmarks/tpch/tpch.py +++ b/benchmarks/tpch/tpch.py @@ -23,7 +23,7 @@ def bench(data_path, query_path) -> None: - with Path.open("results.csv", "w") as results: + with Path("results.csv").open("w") as results: # register tables start = time.time() total_time_millis = 0 @@ -46,7 +46,7 @@ def bench(data_path, query_path) -> None: print("Configuration:\n", ctx) # register tables - with Path.open("create_tables.sql") as f: + with Path("create_tables.sql").open() as f: sql = "" for line in f.readlines(): if line.startswith("--"): @@ -66,7 +66,7 @@ def bench(data_path, query_path) -> None: # run queries for query in range(1, 23): - with Path.open(f"{query_path}/q{query}.sql") as f: + with Path(f"{query_path}/q{query}.sql").open() as f: text = f.read() tmp = text.split(";") queries = [s.strip() for s in tmp if len(s.strip()) > 0] diff --git a/python/tests/test_config.py b/conftest.py similarity index 59% rename from python/tests/test_config.py rename to conftest.py index c1d7f97e1..0c9410636 100644 --- a/python/tests/test_config.py +++ b/conftest.py @@ -15,28 +15,22 @@ # 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" +"""Pytest configuration for doctest namespace injection.""" - 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 +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 + + +@pytest.fixture(autouse=True) +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/Cargo.toml b/crates/core/Cargo.toml new file mode 100644 index 000000000..c5f1e0167 --- /dev/null +++ b/crates/core/Cargo.toml @@ -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. + +[package] +name = "datafusion-python" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description.workspace = true +homepage.workspace = true +repository.workspace = true +include = [ + "src", + "../LICENSE.txt", + "build.rs", + "../pyproject.toml", + "Cargo.toml", + "../Cargo.lock", +] + +[dependencies] +tokio = { workspace = true, features = [ + "macros", + "rt", + "rt-multi-thread", + "sync", +] } +pyo3 = { workspace = true, features = [ + "extension-module", + "generate-import-lib", +] } +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"] } +datafusion-substrait = { workspace = true, optional = true } +datafusion-proto = { workspace = true } +datafusion-ffi = { workspace = true } +datafusion-spark = { workspace = true, features = ["core"] } +prost = { workspace = true } # keep in line with `datafusion-substrait` +serde_json = { workspace = true } +uuid = { workspace = true, features = ["v4"] } +mimalloc = { workspace = true, optional = true, features = [ + "local_dynamic_tls", + # Pin to mimalloc v2 until apache/datafusion-python#1607 resolves. + "v2", +] } +async-trait = { workspace = true } +futures = { workspace = true } +cstr = { workspace = true } +object_store = { workspace = true, features = ["aws", "gcp", "azure", "http"] } +url = { workspace = true } +log = { workspace = true } +parking_lot = { workspace = true } +datafusion-python-util = { workspace = true } + +[build-dependencies] +prost-types = { workspace = true } +pyo3-build-config = { workspace = true } + +[features] +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"] + +[lib] +name = "datafusion_python" +crate-type = ["cdylib", "rlib"] diff --git a/build.rs b/crates/core/build.rs similarity index 100% rename from build.rs rename to crates/core/build.rs 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/array.rs b/crates/core/src/array.rs new file mode 100644 index 000000000..dfe963183 --- /dev/null +++ b/crates/core/src/array.rs @@ -0,0 +1,88 @@ +// 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::ptr::NonNull; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef}; +use arrow::datatypes::{Field, FieldRef}; +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::pyarrow::ToPyArrow; +use pyo3::prelude::{PyAnyMethods, PyCapsuleMethods}; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods}; + +use crate::errors::PyDataFusionResult; + +/// A Python object which implements the Arrow PyCapsule for importing +/// into other libraries. +#[pyclass( + from_py_object, + name = "ArrowArrayExportable", + module = "datafusion", + frozen +)] +#[derive(Clone)] +pub struct PyArrowArrayExportable { + array: ArrayRef, + field: FieldRef, +} + +#[pymethods] +impl PyArrowArrayExportable { + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_array__<'py>( + &'py self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyDataFusionResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> { + let field = if let Some(schema_capsule) = requested_schema { + let data: NonNull = schema_capsule + .pointer_checked(Some(c"arrow_schema"))? + .cast(); + let schema_ptr = unsafe { data.as_ref() }; + let desired_field = Field::try_from(schema_ptr)?; + + Arc::new(desired_field) + } else { + Arc::clone(&self.field) + }; + + let ffi_schema = FFI_ArrowSchema::try_from(&field)?; + let schema_capsule = PyCapsule::new_with_value(py, ffi_schema, cr"arrow_schema")?; + + let ffi_array = FFI_ArrowArray::new(&self.array.to_data()); + let array_capsule = PyCapsule::new_with_value(py, ffi_array, cr"arrow_array")?; + + Ok((schema_capsule, array_capsule)) + } +} + +impl ToPyArrow for PyArrowArrayExportable { + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { + let module = py.import("pyarrow")?; + let method = module.getattr("array")?; + let array = method.call((self.clone(),), None)?; + Ok(array) + } +} + +impl PyArrowArrayExportable { + pub fn new(array: ArrayRef, field: FieldRef) -> Self { + Self { array, field } + } +} diff --git a/crates/core/src/catalog.rs b/crates/core/src/catalog.rs new file mode 100644 index 000000000..8ad49b098 --- /dev/null +++ b/crates/core/src/catalog.rs @@ -0,0 +1,701 @@ +// 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::HashSet; +use std::ptr::NonNull; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::{ + CatalogProvider, CatalogProviderList, MemoryCatalogProvider, MemoryCatalogProviderList, + MemorySchemaProvider, SchemaProvider, +}; +use datafusion::common::DataFusionError; +use datafusion::datasource::TableProvider; +use datafusion_ffi::catalog_provider::FFI_CatalogProvider; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::schema_provider::FFI_SchemaProvider; +use datafusion_python_util::{ + create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, wait_for_future, +}; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyKeyError; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +use crate::context::PySessionContext; +use crate::dataset::Dataset; +use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err}; +use crate::table::PyTable; + +#[pyclass( + from_py_object, + frozen, + name = "RawCatalogList", + module = "datafusion.catalog", + subclass +)] +#[derive(Clone)] +pub struct PyCatalogList { + pub catalog_list: Arc, + codec: Arc, +} + +#[pyclass( + from_py_object, + frozen, + name = "RawCatalog", + module = "datafusion.catalog", + subclass +)] +#[derive(Clone)] +pub struct PyCatalog { + pub catalog: Arc, + codec: Arc, +} + +#[pyclass( + from_py_object, + frozen, + name = "RawSchema", + module = "datafusion.catalog", + subclass +)] +#[derive(Clone)] +pub struct PySchema { + pub schema: Arc, + codec: Arc, +} + +impl PyCatalog { + pub(crate) fn new_from_parts( + catalog: Arc, + codec: Arc, + ) -> Self { + Self { catalog, codec } + } +} + +impl PySchema { + pub(crate) fn new_from_parts( + schema: Arc, + codec: Arc, + ) -> Self { + Self { schema, codec } + } +} + +#[pymethods] +impl PyCatalogList { + #[new] + pub fn new( + py: Python, + catalog_list: Py, + session: Option>, + ) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog_list = Arc::new(RustWrappedPyCatalogProviderList::new( + catalog_list, + codec.clone(), + )) as Arc; + Ok(Self { + catalog_list, + codec, + }) + } + + #[staticmethod] + pub fn memory_catalog_list(py: Python, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog_list = + Arc::new(MemoryCatalogProviderList::default()) as Arc; + Ok(Self { + catalog_list, + codec, + }) + } + + pub fn catalog_names(&self) -> HashSet { + self.catalog_list.catalog_names().into_iter().collect() + } + + #[pyo3(signature = (name="public"))] + pub fn catalog(&self, name: &str) -> PyResult> { + let catalog = self + .catalog_list + .catalog(name) + .ok_or(PyKeyError::new_err(format!( + "Schema with name {name} doesn't exist." + )))?; + + 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<()> { + let provider = extract_catalog_provider_from_pyobj(catalog_provider, self.codec.as_ref())?; + + let _ = self + .catalog_list + .register_catalog(name.to_owned(), provider); + + Ok(()) + } + + pub fn __repr__(&self) -> PyResult { + let mut names: Vec = self.catalog_names().into_iter().collect(); + names.sort(); + Ok(format!("CatalogList(catalog_names=[{}])", names.join(", "))) + } +} + +#[pymethods] +impl PyCatalog { + #[new] + pub fn new(py: Python, catalog: Py, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog = Arc::new(RustWrappedPyCatalogProvider::new(catalog, codec.clone())) + as Arc; + Ok(Self { catalog, codec }) + } + + #[staticmethod] + pub fn memory_catalog(py: Python, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog = Arc::new(MemoryCatalogProvider::default()) as Arc; + Ok(Self { catalog, codec }) + } + + pub fn schema_names(&self) -> HashSet { + self.catalog.schema_names().into_iter().collect() + } + + #[pyo3(signature = (name="public"))] + pub fn schema(&self, name: &str) -> PyResult> { + let schema = self + .catalog + .schema(name) + .ok_or(PyKeyError::new_err(format!( + "Schema with name {name} doesn't exist." + )))?; + + 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<()> { + let provider = extract_schema_provider_from_pyobj(schema_provider, self.codec.as_ref())?; + + let _ = self + .catalog + .register_schema(name, provider) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + pub fn deregister_schema(&self, name: &str, cascade: bool) -> PyResult<()> { + let _ = self + .catalog + .deregister_schema(name, cascade) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + pub fn __repr__(&self) -> PyResult { + let mut names: Vec = self.schema_names().into_iter().collect(); + names.sort(); + Ok(format!("Catalog(schema_names=[{}])", names.join(", "))) + } +} + +#[pymethods] +impl PySchema { + #[new] + pub fn new( + py: Python, + schema_provider: Py, + session: Option>, + ) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let schema = + Arc::new(RustWrappedPySchemaProvider::new(schema_provider)) as Arc; + Ok(Self { schema, codec }) + } + + #[staticmethod] + fn memory_schema(py: Python, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let schema = Arc::new(MemorySchemaProvider::default()) as Arc; + Ok(Self { schema, codec }) + } + + #[getter] + fn table_names(&self) -> HashSet { + self.schema.table_names().into_iter().collect() + } + + fn table(&self, name: &str, py: Python) -> PyDataFusionResult { + if let Some(table) = wait_for_future(py, self.schema.table(name))?? { + Ok(PyTable::from(table)) + } else { + Err(PyDataFusionError::Common(format!( + "Table not found: {name}" + ))) + } + } + + fn __repr__(&self) -> PyResult { + let mut names: Vec = self.table_names().into_iter().collect(); + names.sort(); + Ok(format!("Schema(table_names=[{}])", names.join(";"))) + } + + fn register_table(&self, name: &str, table_provider: Bound<'_, PyAny>) -> PyResult<()> { + let py = table_provider.py(); + let codec_capsule = create_logical_extension_capsule(py, self.codec.as_ref())? + .as_any() + .clone(); + + let table = PyTable::new(table_provider, Some(codec_capsule))?; + + let _ = self + .schema + .register_table(name.to_string(), table.table) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + fn deregister_table(&self, name: &str) -> PyResult<()> { + let _ = self + .schema + .deregister_table(name) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + fn table_exist(&self, name: &str) -> bool { + self.schema.table_exist(name) + } +} + +#[derive(Debug)] +pub(crate) struct RustWrappedPySchemaProvider { + schema_provider: Py, + owner_name: Option, +} + +impl RustWrappedPySchemaProvider { + pub fn new(schema_provider: Py) -> Self { + let owner_name = Python::attach(|py| { + schema_provider + .bind(py) + .getattr("owner_name") + .ok() + .map(|name| name.to_string()) + }); + + Self { + schema_provider, + owner_name, + } + } + + fn table_inner(&self, name: &str) -> PyResult>> { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + let py_table_method = provider.getattr("table")?; + + let py_table = py_table_method.call((name,), None)?; + if py_table.is_none() { + return Ok(None); + } + + let table = PyTable::new(py_table, None)?; + + Ok(Some(table.table)) + }) + } +} + +#[async_trait] +impl SchemaProvider for RustWrappedPySchemaProvider { + fn owner_name(&self) -> Option<&str> { + self.owner_name.as_deref() + } + + fn table_names(&self) -> Vec { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + + provider + .getattr("table_names") + .and_then(|names| names.extract::>()) + .unwrap_or_else(|err| { + log::error!("Unable to get table_names: {err}"); + Vec::default() + }) + }) + } + + async fn table( + &self, + name: &str, + ) -> datafusion::common::Result>, DataFusionError> { + self.table_inner(name) + .map_err(|e| DataFusionError::External(Box::new(e))) + } + + fn register_table( + &self, + name: String, + table: Arc, + ) -> datafusion::common::Result>> { + let py_table = PyTable::from(table); + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + let _ = provider + .call_method1("register_table", (name, py_table)) + .map_err(to_datafusion_err)?; + // Since the definition of `register_table` says that an error + // will be returned if the table already exists, there is no + // case where we want to return a table provider as output. + Ok(None) + }) + } + + fn deregister_table( + &self, + name: &str, + ) -> datafusion::common::Result>> { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + let table = provider + .call_method1("deregister_table", (name,)) + .map_err(to_datafusion_err)?; + if table.is_none() { + return Ok(None); + } + + // If we can turn this table provider into a `Dataset`, return it. + // Otherwise, return None. + let dataset = match Dataset::new(&table, py) { + Ok(dataset) => Some(Arc::new(dataset) as Arc), + Err(_) => None, + }; + + Ok(dataset) + }) + } + + fn table_exist(&self, name: &str) -> bool { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + provider + .call_method1("table_exist", (name,)) + .and_then(|pyobj| pyobj.extract()) + .unwrap_or(false) + }) + } +} + +#[derive(Debug)] +pub(crate) struct RustWrappedPyCatalogProvider { + pub(crate) catalog_provider: Py, + codec: Arc, +} + +impl RustWrappedPyCatalogProvider { + pub fn new(catalog_provider: Py, codec: Arc) -> Self { + Self { + catalog_provider, + codec, + } + } + + fn schema_inner(&self, name: &str) -> PyResult>> { + Python::attach(|py| { + let provider = self.catalog_provider.bind(py); + + let py_schema = provider.call_method1("schema", (name,))?; + if py_schema.is_none() { + return Ok(None); + } + + extract_schema_provider_from_pyobj(py_schema, self.codec.as_ref()).map(Some) + }) + } +} + +#[async_trait] +impl CatalogProvider for RustWrappedPyCatalogProvider { + fn schema_names(&self) -> Vec { + Python::attach(|py| { + let provider = self.catalog_provider.bind(py); + provider + .call_method0("schema_names") + .and_then(|names| names.extract::>()) + .map(|names| names.into_iter().collect()) + .unwrap_or_else(|err| { + log::error!("Unable to get schema_names: {err}"); + Vec::default() + }) + }) + } + + fn schema(&self, name: &str) -> Option> { + self.schema_inner(name).unwrap_or_else(|err| { + log::error!("CatalogProvider schema returned error: {err}"); + None + }) + } + + fn register_schema( + &self, + name: &str, + schema: Arc, + ) -> datafusion::common::Result>> { + Python::attach(|py| { + 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) + .map_err(to_datafusion_err)?, + }; + + let provider = self.catalog_provider.bind(py); + let schema = provider + .call_method1("register_schema", (name, py_schema)) + .map_err(to_datafusion_err)?; + if schema.is_none() { + return Ok(None); + } + + let schema = Arc::new(RustWrappedPySchemaProvider::new(schema.into())) + as Arc; + + Ok(Some(schema)) + }) + } + + fn deregister_schema( + &self, + name: &str, + cascade: bool, + ) -> datafusion::common::Result>> { + Python::attach(|py| { + let provider = self.catalog_provider.bind(py); + let schema = provider + .call_method1("deregister_schema", (name, cascade)) + .map_err(to_datafusion_err)?; + if schema.is_none() { + return Ok(None); + } + + let schema = Arc::new(RustWrappedPySchemaProvider::new(schema.into())) + as Arc; + + Ok(Some(schema)) + }) + } +} + +#[derive(Debug)] +pub(crate) struct RustWrappedPyCatalogProviderList { + pub(crate) catalog_provider_list: Py, + codec: Arc, +} + +impl RustWrappedPyCatalogProviderList { + pub fn new(catalog_provider_list: Py, codec: Arc) -> Self { + Self { + catalog_provider_list, + codec, + } + } + + fn catalog_inner(&self, name: &str) -> PyResult>> { + Python::attach(|py| { + let provider = self.catalog_provider_list.bind(py); + + let py_schema = provider.call_method1("catalog", (name,))?; + if py_schema.is_none() { + return Ok(None); + } + + extract_catalog_provider_from_pyobj(py_schema, self.codec.as_ref()).map(Some) + }) + } +} + +#[async_trait] +impl CatalogProviderList for RustWrappedPyCatalogProviderList { + fn catalog_names(&self) -> Vec { + Python::attach(|py| { + let provider = self.catalog_provider_list.bind(py); + provider + .call_method0("catalog_names") + .and_then(|names| names.extract::>()) + .map(|names| names.into_iter().collect()) + .unwrap_or_else(|err| { + log::error!("Unable to get catalog_names: {err}"); + Vec::default() + }) + }) + } + + fn catalog(&self, name: &str) -> Option> { + self.catalog_inner(name).unwrap_or_else(|err| { + log::error!("CatalogProvider catalog returned error: {err}"); + None + }) + } + + fn register_catalog( + &self, + name: String, + catalog: Arc, + ) -> Option> { + Python::attach(|py| { + 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) { + Ok(c) => c, + Err(err) => { + log::error!( + "register_catalog returned error during conversion to PyAny: {err}" + ); + return None; + } + } + } + }; + + let provider = self.catalog_provider_list.bind(py); + let catalog = match provider.call_method1("register_catalog", (name, py_catalog)) { + Ok(c) => c, + Err(err) => { + log::error!("register_catalog returned error: {err}"); + return None; + } + }; + if catalog.is_none() { + return None; + } + + let catalog = Arc::new(RustWrappedPyCatalogProvider::new( + catalog.into(), + self.codec.clone(), + )) as Arc; + + Some(catalog) + }) + } +} + +fn extract_catalog_provider_from_pyobj( + mut catalog_provider: Bound, + codec: &FFI_LogicalExtensionCodec, +) -> PyResult> { + if catalog_provider.hasattr("__datafusion_catalog_provider__")? { + let py = catalog_provider.py(); + let codec_capsule = create_logical_extension_capsule(py, codec)?; + catalog_provider = catalog_provider + .getattr("__datafusion_catalog_provider__")? + .call1((codec_capsule,))?; + } + + let provider = if let Ok(capsule) = catalog_provider.cast::() { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_catalog_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider + } else { + match catalog_provider.extract::() { + Ok(py_catalog) => py_catalog.catalog, + Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( + catalog_provider.into(), + Arc::new(codec.clone()), + )) as Arc, + } + }; + + Ok(provider) +} + +fn extract_schema_provider_from_pyobj( + mut schema_provider: Bound, + codec: &FFI_LogicalExtensionCodec, +) -> PyResult> { + if schema_provider.hasattr("__datafusion_schema_provider__")? { + let py = schema_provider.py(); + let codec_capsule = create_logical_extension_capsule(py, codec)?; + schema_provider = schema_provider + .getattr("__datafusion_schema_provider__")? + .call1((codec_capsule,))?; + } + + let provider = if let Ok(capsule) = schema_provider.cast::() { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_schema_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider + } else { + match schema_provider.extract::() { + Ok(py_schema) => py_schema.schema, + Err(_) => Arc::new(RustWrappedPySchemaProvider::new(schema_provider.into())) + as Arc, + } + }; + + Ok(provider) +} + +fn extract_logical_extension_codec( + py: Python, + obj: Option>, +) -> PyResult> { + let obj = match obj { + Some(obj) => obj, + None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, + }; + ffi_logical_codec_from_pycapsule(obj).map(Arc::new) +} + +pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + Ok(()) +} diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs new file mode 100644 index 000000000..26853e69f --- /dev/null +++ b/crates/core/src/codec.rs @@ -0,0 +1,1164 @@ +// 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 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 family registry +//! +//! | 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) | +//! +//! 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::{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, AggregateUDFImpl, Extension, LogicalPlan, ScalarUDF, ScalarUDFImpl, Signature, + TypeSignature, Volatility, WindowUDF, WindowUDFImpl, +}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; +use datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, +}; +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. +// +// 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"; + +/// 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; + +/// 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 +/// (`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, + python_udf_inlining: bool, +} + +impl PythonLogicalCodec { + pub fn new(inner: Arc) -> Self { + 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. 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 + } + + pub fn python_udf_inlining(&self) -> bool { + self.python_udf_inlining + } +} + +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<()> { + 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 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 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 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 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 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`) +/// 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_FAMILY`] et al.) so the wire format is identical. +#[derive(Debug)] +pub struct PythonPhysicalCodec { + inner: Arc, + python_udf_inlining: bool, +} + +impl PythonPhysicalCodec { + pub fn new(inner: Arc) -> Self { + 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 { + fn default() -> Self { + Self::new(Arc::new(DefaultPhysicalExtensionCodec {})) + } +} + +impl PhysicalExtensionCodec for PythonPhysicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + self.inner.try_encode(node, buf, proto_converter) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + 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 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_expr( + &self, + node: &Arc, + buf: &mut Vec, + ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result<()> { + self.inner.try_encode_expr(node, buf, ctx) + } + + fn try_decode_expr( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalExprDecodeCtx<'_>, + ) -> Result> { + self.inner.try_decode_expr(buf, inputs, ctx) + } + + fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { + 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 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 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 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) + } +} + +// ============================================================================= +// 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 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) + }) +} + +/// 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>> { + 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 { + return Ok(None); + }; + let udf = decode_python_scalar_udf(py, payload).map_err(to_datafusion_err)?; + 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) +} + +/// 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> { + 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)) +} + +/// 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 +/// 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()) +} + +// ============================================================================= +// 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>> { + 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 { + 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>> { + 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 { + 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::*; + + 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_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"); + + let payload = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY) + .unwrap() + .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/src/common.rs b/crates/core/src/common.rs similarity index 100% rename from src/common.rs rename to crates/core/src/common.rs diff --git a/src/common/data_type.rs b/crates/core/src/common/data_type.rs similarity index 96% rename from src/common/data_type.rs rename to crates/core/src/common/data_type.rs index 55848da5c..e79aea4ef 100644 --- a/src/common/data_type.rs +++ b/crates/core/src/common/data_type.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use datafusion::arrow::array::Array; use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit}; use datafusion::common::ScalarValue; @@ -22,6 +24,9 @@ use datafusion::logical_expr::expr::NullTreatment as DFNullTreatment; use pyo3::exceptions::{PyNotImplementedError, PyValueError}; use pyo3::prelude::*; +/// A [`ScalarValue`] wrapped in a Python object. This struct allows for conversion +/// from a variety of Python objects into a [`ScalarValue`]. See +/// ``FromPyArrow::from_pyarrow_bound`` conversion details. #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] pub struct PyScalarValue(pub ScalarValue); @@ -37,7 +42,14 @@ impl From for ScalarValue { } #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(frozen, eq, eq_int, name = "RexType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "RexType", + module = "datafusion.common" +)] pub enum RexType { Alias, Literal, @@ -58,7 +70,12 @@ pub enum RexType { /// to map types from one system to another. // TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now #[derive(Debug, Clone)] -#[pyclass(name = "DataTypeMap", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + name = "DataTypeMap", + module = "datafusion.common", + subclass +)] pub struct DataTypeMap { #[pyo3(get, set)] pub arrow_type: PyDataType, @@ -317,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(), )), @@ -329,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)), @@ -344,6 +361,10 @@ impl DataTypeMap { ScalarValue::Map(_) => Err(PyNotImplementedError::new_err( "ScalarValue::Map".to_string(), )), + ScalarValue::RunEndEncoded(field1, field2, _) => Ok(DataType::RunEndEncoded( + Arc::clone(field1), + Arc::clone(field2), + )), } } } @@ -584,7 +605,12 @@ impl DataTypeMap { /// Since `DataType` exists in another package we cannot make that happen here so we wrap /// `DataType` as `PyDataType` This exists solely to satisfy those constraints. #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(frozen, name = "DataType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + name = "DataType", + module = "datafusion.common" +)] pub struct PyDataType { pub data_type: DataType, } @@ -642,7 +668,14 @@ impl From for PyDataType { /// Represents the possible Python types that can be mapped to the SQL types #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(frozen, eq, eq_int, name = "PythonType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "PythonType", + module = "datafusion.common" +)] pub enum PythonType { Array, Bool, @@ -662,7 +695,14 @@ pub enum PythonType { #[allow(non_camel_case_types)] #[allow(clippy::upper_case_acronyms)] #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(frozen, eq, eq_int, name = "SqlType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "SqlType", + module = "datafusion.common" +)] pub enum SqlType { ANY, ARRAY, @@ -721,6 +761,7 @@ pub enum SqlType { #[allow(clippy::upper_case_acronyms)] #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] #[pyclass( + from_py_object, frozen, eq, eq_int, diff --git a/src/common/df_schema.rs b/crates/core/src/common/df_schema.rs similarity index 93% rename from src/common/df_schema.rs rename to crates/core/src/common/df_schema.rs index eb62469cf..9167e772e 100644 --- a/src/common/df_schema.rs +++ b/crates/core/src/common/df_schema.rs @@ -21,7 +21,13 @@ use datafusion::common::DFSchema; use pyo3::prelude::*; #[derive(Debug, Clone)] -#[pyclass(frozen, name = "DFSchema", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DFSchema", + module = "datafusion.common", + subclass +)] pub struct PyDFSchema { schema: Arc, } diff --git a/src/common/function.rs b/crates/core/src/common/function.rs similarity index 93% rename from src/common/function.rs rename to crates/core/src/common/function.rs index bc6f23160..41cab515f 100644 --- a/src/common/function.rs +++ b/crates/core/src/common/function.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use super::data_type::PyDataType; -#[pyclass(frozen, name = "SqlFunction", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SqlFunction", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlFunction { pub name: String, diff --git a/src/common/schema.rs b/crates/core/src/common/schema.rs similarity index 92% rename from src/common/schema.rs rename to crates/core/src/common/schema.rs index 4e46592aa..94b3ce0ae 100644 --- a/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; @@ -34,7 +33,13 @@ use super::data_type::DataTypeMap; use super::function::SqlFunction; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "SqlSchema", module = "datafusion.common", subclass, frozen)] +#[pyclass( + from_py_object, + name = "SqlSchema", + module = "datafusion.common", + subclass, + frozen +)] #[derive(Debug, Clone)] pub struct SqlSchema { name: Arc>, @@ -43,7 +48,12 @@ pub struct SqlSchema { functions: Arc>>, } -#[pyclass(name = "SqlTable", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + name = "SqlTable", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlTable { #[pyo3(get, set)] @@ -87,7 +97,12 @@ impl SqlTable { } } -#[pyclass(name = "SqlView", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + name = "SqlView", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlView { #[pyo3(get, set)] @@ -203,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() } @@ -247,7 +258,13 @@ fn is_supported_push_down_expr(_expr: &Expr) -> bool { true } -#[pyclass(frozen, name = "SqlStatistics", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SqlStatistics", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlStatistics { row_count: f64, @@ -266,7 +283,13 @@ impl SqlStatistics { } } -#[pyclass(frozen, name = "Constraints", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Constraints", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyConstraints { pub constraints: Constraints, @@ -291,7 +314,14 @@ impl Display for PyConstraints { } #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(frozen, eq, eq_int, name = "TableType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "TableType", + module = "datafusion.common" +)] pub enum PyTableType { Base, View, @@ -318,7 +348,13 @@ impl From for PyTableType { } } -#[pyclass(frozen, name = "TableSource", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "TableSource", + module = "datafusion.common", + subclass +)] #[derive(Clone)] pub struct PyTableSource { pub table_source: Arc, diff --git a/src/context.rs b/crates/core/src/context.rs similarity index 52% rename from src/context.rs rename to crates/core/src/context.rs index ad4fc36b1..7bbeed2f1 100644 --- a/src/context.rs +++ b/crates/core/src/context.rs @@ -16,7 +16,8 @@ // 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; @@ -26,8 +27,8 @@ use arrow::pyarrow::FromPyArrow; use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; -use datafusion::catalog::CatalogProvider; -use datafusion::common::{exec_err, ScalarValue, TableReference}; +use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory}; +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::{ @@ -39,42 +40,69 @@ 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::execution::{FunctionRegistry, TaskContextProvider}; use datafusion::prelude::{ - AvroReadOptions, CsvReadOptions, DataFrame, NdJsonReadOptions, ParquetReadOptions, + AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions, +}; +use datafusion_ffi::catalog_provider::FFI_CatalogProvider; +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::LogicalExtensionCodec; +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, physical_optimizer_rule_from_pycapsule, spawn_future, + wait_for_future, }; -use datafusion_ffi::catalog_provider::{FFI_CatalogProvider, ForeignCatalogProvider}; use object_store::ObjectStore; -use pyo3::exceptions::{PyKeyError, PyValueError}; -use pyo3::prelude::*; -use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple, PyType}; use pyo3::IntoPyObjectExt; +use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; use url::Url; use uuid::Uuid; -use crate::catalog::{PyCatalog, RustWrappedPyCatalogProvider}; +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; use crate::dataset::Dataset; -use crate::errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult}; +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; use crate::record_batch::PyRecordBatchStream; -use crate::sql::exceptions::py_value_err; use crate::sql::logical::PyLogicalPlan; use crate::sql::util::replace_placeholders_with_strings; use crate::store::StorageContexts; -use crate::table::PyTable; +use crate::table::{PyTable, RustWrappedPyTableProviderFactory}; use crate::udaf::PyAggregateUDF; use crate::udf::PyScalarUDF; use crate::udtf::PyTableFunction; use crate::udwf::PyWindowUDF; -use crate::utils::{get_global_ctx, spawn_future, validate_pycapsule, wait_for_future}; /// Configuration options for a SessionContext -#[pyclass(frozen, name = "SessionConfig", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SessionConfig", + module = "datafusion", + subclass +)] #[derive(Clone, Default)] pub struct PySessionConfig { pub config: SessionConfig, @@ -164,10 +192,43 @@ impl PySessionConfig { fn set(&self, key: &str, value: &str) -> Self { Self::from(self.config.clone().set_str(key, value)) } + + pub fn with_extension(&self, extension: Bound) -> PyResult { + if !extension.hasattr("__datafusion_extension_options__")? { + return Err(pyo3::exceptions::PyAttributeError::new_err( + "Expected extension object to define __datafusion_extension_options__()", + )); + } + let capsule = extension.call_method0("__datafusion_extension_options__")?; + let capsule = capsule.cast::()?; + + let extension: NonNull = capsule + .pointer_checked(Some(c"datafusion_extension_options"))? + .cast(); + let mut extension = unsafe { extension.as_ref() }.clone(); + + let mut config = self.config.clone(); + let options = config.options_mut(); + if let Some(prior_extension) = options.extensions.get::() { + extension + .merge(prior_extension) + .map_err(py_datafusion_err)?; + } + + options.extensions.insert(extension); + + Ok(Self::from(config)) + } } /// Runtime options for a SessionContext -#[pyclass(frozen, name = "RuntimeEnvBuilder", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "RuntimeEnvBuilder", + module = "datafusion", + subclass +)] #[derive(Clone)] pub struct PyRuntimeEnvBuilder { pub builder: RuntimeEnvBuilder, @@ -254,7 +315,13 @@ impl PyRuntimeEnvBuilder { } /// `PySQLOptions` allows you to specify options to the sql execution. -#[pyclass(frozen, name = "SQLOptions", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SQLOptions", + module = "datafusion", + subclass +)] #[derive(Clone)] pub struct PySQLOptions { pub options: SQLOptions, @@ -293,10 +360,18 @@ impl PySQLOptions { /// `PySessionContext` is able to plan and execute DataFusion plans. /// It has a powerful optimizer, a physical planner for local execution, and a /// multi-threaded execution engine to perform the execution. -#[pyclass(frozen, name = "SessionContext", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SessionContext", + module = "datafusion", + subclass +)] #[derive(Clone)] pub struct PySessionContext { - pub ctx: SessionContext, + pub ctx: Arc, + logical_codec: Arc, + physical_codec: Arc, } #[pymethods] @@ -322,23 +397,32 @@ 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 { - ctx: SessionContext::new_with_state(session_state), + ctx, + logical_codec: Arc::new(PythonLogicalCodec::default()), + physical_codec: Arc::new(PythonPhysicalCodec::default()), }) } pub fn enable_url_table(&self) -> PyResult { Ok(PySessionContext { - ctx: self.ctx.clone().enable_url_table(), + 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), }) } - #[classmethod] + #[staticmethod] #[pyo3(signature = ())] - fn global_ctx(_cls: &Bound<'_, PyType>) -> PyResult { + pub fn global_ctx() -> PyResult { + let ctx = get_global_ctx().clone(); Ok(Self { - ctx: get_global_ctx().clone(), + ctx, + logical_codec: Arc::new(PythonLogicalCodec::default()), + physical_codec: Arc::new(PythonPhysicalCodec::default()), }) } @@ -366,11 +450,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", @@ -379,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>, @@ -388,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 => { @@ -424,6 +511,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, @@ -453,7 +544,8 @@ impl PySessionContext { let mut df = wait_for_future(py, async { self.ctx.sql_with_options(&query, options).await - })??; + })? + .map_err(from_datafusion_error)?; if !param_values.is_empty() { df = df.with_param_values(param_values)?; @@ -606,7 +698,8 @@ impl PySessionContext { } pub fn register_table(&self, name: &str, table: Bound<'_, PyAny>) -> PyDataFusionResult<()> { - let table = PyTable::new(&table)?; + let session = self.clone().into_bound_py_any(table.py())?; + let table = PyTable::new(table, Some(session))?; self.ctx.register_table(name, table.table)?; Ok(()) @@ -617,26 +710,105 @@ impl PySessionContext { Ok(()) } + pub fn register_table_factory( + &self, + format: &str, + mut factory: Bound<'_, PyAny>, + ) -> PyDataFusionResult<()> { + if factory.hasattr("__datafusion_table_provider_factory__")? { + let py = factory.py(); + 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,))?; + } + + let factory: Arc = + if let Ok(capsule) = factory.cast::().map_err(py_datafusion_err) { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_table_provider_factory"))? + .cast(); + let factory = unsafe { data.as_ref() }; + factory.into() + } else { + Arc::new(RustWrappedPyTableProviderFactory::new( + factory.into(), + self.ffi_logical_codec(), + )) + }; + + let st = self.ctx.state_ref(); + let mut lock = st.write(); + lock.table_factories_mut() + .insert(format.to_owned(), factory); + + Ok(()) + } + + pub fn register_catalog_provider_list( + &self, + mut provider: Bound, + ) -> PyDataFusionResult<()> { + if provider.hasattr("__datafusion_catalog_provider_list__")? { + let py = provider.py(); + 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,))?; + } + + let provider = if let Ok(capsule) = provider.cast::() { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_catalog_provider_list"))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider + } else { + match provider.extract::() { + Ok(py_catalog_list) => py_catalog_list.catalog_list, + Err(_) => Arc::new(RustWrappedPyCatalogProviderList::new( + provider.into(), + self.ffi_logical_codec(), + )) as Arc, + } + }; + + self.ctx.register_catalog_list(provider); + + Ok(()) + } + pub fn register_catalog_provider( &self, name: &str, - provider: Bound<'_, PyAny>, + mut provider: Bound<'_, PyAny>, ) -> PyDataFusionResult<()> { - let provider = if provider.hasattr("__datafusion_catalog_provider__")? { - let capsule = provider + if provider.hasattr("__datafusion_catalog_provider__")? { + let py = provider.py(); + let ffi = self.ffi_logical_codec(); + let codec_capsule = create_logical_extension_capsule(py, ffi.as_ref())?; + provider = provider .getattr("__datafusion_catalog_provider__")? - .call0()?; - let capsule = capsule.downcast::().map_err(py_datafusion_err)?; - validate_pycapsule(capsule, "datafusion_catalog_provider")?; + .call1((codec_capsule,))?; + } - let provider = unsafe { capsule.reference::() }; - let provider: ForeignCatalogProvider = provider.into(); - Arc::new(provider) as Arc + let provider = if let Ok(capsule) = provider.cast::() { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_catalog_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider } else { match provider.extract::() { Ok(py_catalog) => py_catalog.catalog, - Err(_) => Arc::new(RustWrappedPyCatalogProvider::new(provider.into())) - as Arc, + Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( + provider.into(), + self.ffi_logical_codec(), + )) as Arc, } }; @@ -660,12 +832,31 @@ impl PySessionContext { name: &str, partitions: PyArrowType>>, ) -> PyDataFusionResult<()> { - let schema = partitions.0[0][0].schema(); + // Take the schema from the first available batch; error instead of + // panicking when the partitions hold no batches. + let schema = partitions + .0 + .iter() + .find_map(|partition| partition.first()) + .ok_or_else(|| { + PyValueError::new_err( + "Cannot register record batches without a schema: the \ + provided partitions contain no record batches.", + ) + })? + .schema(); let table = MemTable::try_new(schema, partitions.0)?; self.ctx.register_table(name, Arc::new(table))?; 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, @@ -676,7 +867,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, @@ -685,72 +876,50 @@ 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(()) } - #[allow(clippy::too_many_arguments)] #[pyo3(signature = (name, path, - schema=None, - has_header=true, - delimiter=",", - schema_infer_max_records=1000, - file_extension=".csv", - file_compression_type=None))] + options=None))] pub fn register_csv( &self, name: &str, path: &Bound<'_, PyAny>, - schema: Option>, - has_header: bool, - delimiter: &str, - schema_infer_max_records: usize, - file_extension: &str, - file_compression_type: Option, + options: Option<&PyCsvReadOptions>, py: Python, ) -> PyDataFusionResult<()> { - let delimiter = delimiter.as_bytes(); - if delimiter.len() != 1 { - return Err(PyDataFusionError::PythonError(py_value_err( - "Delimiter must be a single character", - ))); - } - - let mut options = CsvReadOptions::new() - .has_header(has_header) - .delimiter(delimiter[0]) - .schema_infer_max_records(schema_infer_max_records) - .file_extension(file_extension) - .file_compression_type(parse_file_compression_type(file_compression_type)?); - options.schema = schema.as_ref().map(|x| &x.0); + 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(()) @@ -775,25 +944,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 = NdJsonReadOptions::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(()) } @@ -812,22 +973,38 @@ 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 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(()) + } - let result = self.ctx.register_avro(name, path, options); - wait_for_future(py, result)??; + #[pyo3(signature = (name, path, schema=None, file_extension=".arrow", table_partition_cols=vec![]))] + pub fn register_arrow( + &self, + name: &str, + path: PathBuf, + schema: Option>, + file_extension: &str, + table_partition_cols: Vec<(String, PyArrowType)>, + py: Python, + ) -> PyDataFusionResult<()> { + 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(()) + } + pub fn register_batch( + &self, + name: &str, + batch: PyArrowType, + ) -> PyDataFusionResult<()> { + self.ctx.register_batch(name, batch.0)?; Ok(()) } @@ -850,62 +1027,113 @@ 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(()) } + /// Register all `datafusion-spark` UDFs/UDAFs/UDWFs, overriding any built-in + /// DataFusion functions of the same name with their Spark-semantics version. + pub fn enable_spark_functions(&self) -> PyResult<()> { + for udf in datafusion_spark::all_default_scalar_functions() { + self.ctx.register_udf((*udf).clone()); + } + for udaf in datafusion_spark::all_default_aggregate_functions() { + self.ctx.register_udaf((*udaf).clone()); + } + for udwf in datafusion_spark::all_default_window_functions() { + self.ctx.register_udwf((*udwf).clone()); + } + 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); + } + + 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, name: &str) -> PyResult> { + pub fn catalog(&self, py: Python, name: &str) -> PyResult> { let catalog = self.ctx.catalog(name).ok_or(PyKeyError::new_err(format!( "Catalog with name {name} doesn't exist." )))?; - Python::attach(|py| { - match catalog - .as_any() - .downcast_ref::() - { - Some(wrapped_schema) => Ok(wrapped_schema.catalog_provider.clone_ref(py)), - None => PyCatalog::from(catalog).into_py_any(py), + 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)?) } - }) + } } pub fn catalog_names(&self) -> HashSet { 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()))?; match res { Ok(df) => Ok(PyDataFrame::new(df)), Err(e) => { - if let datafusion::error::DataFusionError::Plan(msg) = &e { - if msg.contains("No table named") { - return Err(PyKeyError::new_err(msg.to_string())); - } + if let datafusion::error::DataFusionError::Plan(msg) = &e + && msg.contains("No table named") + { + return Err(PyKeyError::new_err(msg.to_string())); } Err(py_datafusion_err(e)) } @@ -924,6 +1152,74 @@ 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() + } + + 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 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 + .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( @@ -936,85 +1232,38 @@ 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 = NdJsonReadOptions::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)) } - #[allow(clippy::too_many_arguments)] #[pyo3(signature = ( path, - schema=None, - has_header=true, - delimiter=",", - schema_infer_max_records=1000, - file_extension=".csv", - table_partition_cols=vec![], - file_compression_type=None))] + options=None))] pub fn read_csv( &self, path: &Bound<'_, PyAny>, - schema: Option>, - has_header: bool, - delimiter: &str, - schema_infer_max_records: usize, - file_extension: &str, - table_partition_cols: Vec<(String, PyArrowType)>, - file_compression_type: Option, + options: Option<&PyCsvReadOptions>, py: Python, ) -> PyDataFusionResult { - let delimiter = delimiter.as_bytes(); - if delimiter.len() != 1 { - return Err(PyDataFusionError::PythonError(py_value_err( - "Delimiter must be a single character", - ))); - }; + let options = convert_csv_options(options)?; - let mut options = CsvReadOptions::new() - .has_header(has_header) - .delimiter(delimiter[0]) - .schema_infer_max_records(schema_infer_max_records) - .file_extension(file_extension) - .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 = schema.as_ref().map(|x| &x.0); - - 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)] @@ -1028,7 +1277,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, @@ -1037,25 +1286,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) } @@ -1063,32 +1305,34 @@ 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: PathBuf, + schema: Option>, + file_extension: &str, + table_partition_cols: Vec<(String, PyArrowType)>, + py: Python, + ) -> PyDataFusionResult { + 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)) } pub fn read_table(&self, table: Bound<'_, PyAny>) -> PyDataFusionResult { - let table = PyTable::new(&table)?; + let session = self.clone().into_bound_py_any(table.py())?; + let table = PyTable::new(table, Some(session))?; let df = self.ctx.read_table(table.table())?; Ok(PyDataFrame::new(df)) } @@ -1122,6 +1366,77 @@ impl PySessionContext { let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?; Ok(PyRecordBatchStream::new(stream)) } + + pub fn __datafusion_task_context_provider__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let ctx_provider = Arc::clone(&self.ctx) as Arc; + let ffi_ctx_provider = FFI_TaskContextProvider::from(&ctx_provider); + + PyCapsule::new_with_value(py, ffi_ctx_provider, cr"datafusion_task_context_provider") + } + + pub fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + 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 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), + }) + } + + 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, + }) + } + + 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 { @@ -1149,7 +1464,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}'" ); @@ -1168,25 +1483,160 @@ impl PySessionContext { .register_table(TableReference::Bare { table: name.into() }, Arc::new(table))?; Ok(()) } + + /// 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(&self.ctx) as Arc; + Arc::new(FFI_LogicalExtensionCodec::new( + 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, + )) + } } pub fn parse_file_compression_type( file_compression_type: Option, ) -> Result { - FileCompressionType::from_str(&*file_compression_type.unwrap_or("".to_string()).as_str()) - .map_err(|_| { - PyValueError::new_err("file_compression_type must one of: gzip, bz2, xz, zstd") - }) + FileCompressionType::from_str(&file_compression_type.unwrap_or_default()).map_err(|_| { + PyValueError::new_err("file_compression_type must be one of: gzip, bz2, xz, zstd") + }) +} + +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 + ctx.ctx.as_ref().clone() } } impl From for PySessionContext { fn from(ctx: SessionContext) -> PySessionContext { - PySessionContext { ctx } + PySessionContext { + ctx: Arc::new(ctx), + logical_codec: Arc::new(PythonLogicalCodec::default()), + physical_codec: Arc::new(PythonPhysicalCodec::default()), + } } } diff --git a/src/dataframe.rs b/crates/core/src/dataframe.rs similarity index 79% rename from src/dataframe.rs rename to crates/core/src/dataframe.rs index d920df71e..b1f305551 100644 --- a/src/dataframe.rs +++ b/crates/core/src/dataframe.rs @@ -16,10 +16,12 @@ // under the License. use std::collections::HashMap; -use std::ffi::{CStr, CString}; +use std::ffi::CStr; +use std::ptr::NonNull; +use std::str::FromStr; use std::sync::Arc; -use arrow::array::{new_null_array, Array, ArrayRef, RecordBatch, RecordBatchReader}; +use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchReader, new_null_array}; use arrow::compute::can_cast_types; use arrow::error::ArrowError; use arrow::ffi::FFI_ArrowSchema; @@ -35,28 +37,33 @@ 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::dml::InsertOp; -use datafusion::logical_expr::SortExpr; +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, + 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}; use parking_lot::Mutex; +use pyo3::PyErr; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods}; -use pyo3::PyErr; -use crate::errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult}; -use crate::expr::sort_expr::{to_sort_expressions, PySortExpr}; +use crate::common::data_type::PyScalarValue; +use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err}; use crate::expr::PyExpr; +use crate::expr::sort_expr::{PySortExpr, to_sort_expressions}; use crate::physical_plan::PyExecutionPlan; -use crate::record_batch::{poll_next_batch, PyRecordBatchStream}; +use crate::record_batch::{PyRecordBatchStream, poll_next_batch}; use crate::sql::logical::PyLogicalPlan; use crate::table::{PyTable, TempViewTable}; -use crate::utils::{ - is_ipython_env, py_obj_to_scalar_value, spawn_future, validate_pycapsule, wait_for_future, -}; /// File-level static CStr for the Arrow array stream capsule name. static ARROW_ARRAY_STREAM_NAME: &CStr = cstr!("arrow_array_stream"); @@ -71,18 +78,18 @@ type SharedCachedBatches = Arc>; pub struct FormatterConfig { /// Maximum memory in bytes to use for display (default: 2MB) pub max_bytes: usize, - /// Minimum number of rows to display (default: 20) + /// Minimum number of rows to display (default: 10) pub min_rows: usize, - /// Number of rows to include in __repr__ output (default: 10) - pub repr_rows: usize, + /// Maximum number of rows to include in __repr__ output (default: 10) + pub max_rows: usize, } impl Default for FormatterConfig { fn default() -> Self { Self { max_bytes: 2 * 1024 * 1024, // 2MB - min_rows: 20, - repr_rows: 10, + min_rows: 10, + max_rows: 10, } } } @@ -102,8 +109,12 @@ impl FormatterConfig { return Err("min_rows must be a positive integer".to_string()); } - if self.repr_rows == 0 { - return Err("repr_rows must be a positive integer".to_string()); + if self.max_rows == 0 { + return Err("max_rows must be a positive integer".to_string()); + } + + if self.min_rows > self.max_rows { + return Err("min_rows must be less than or equal to max_rows".to_string()); } Ok(()) @@ -135,11 +146,11 @@ fn import_python_formatter(py: Python<'_>) -> PyResult> { // Helper function to extract attributes with fallback to default fn get_attr<'a, T>(py_object: &'a Bound<'a, PyAny>, attr_name: &str, default_value: T) -> T where - T: for<'py> pyo3::FromPyObject<'py> + Clone, + T: for<'py> pyo3::FromPyObject<'py, 'py> + Clone, { py_object .getattr(attr_name) - .and_then(|v| v.extract::()) + .and_then(|v| v.extract::().map_err(Into::::into)) .unwrap_or_else(|_| default_value.clone()) } @@ -147,13 +158,30 @@ where fn build_formatter_config_from_python(formatter: &Bound<'_, PyAny>) -> PyResult { let default_config = FormatterConfig::default(); let max_bytes = get_attr(formatter, "max_memory_bytes", default_config.max_bytes); - let min_rows = get_attr(formatter, "min_rows_display", default_config.min_rows); - let repr_rows = get_attr(formatter, "repr_rows", default_config.repr_rows); + let min_rows = get_attr(formatter, "min_rows", default_config.min_rows); + + // Backward compatibility: Try max_rows first (new name), fall back to repr_rows (deprecated), + // then use default. This ensures backward compatibility with custom formatter implementations + // during the deprecation period. + let max_rows = get_attr(formatter, "max_rows", 0usize); + let max_rows = if max_rows > 0 { + // max_rows attribute exists and has a value + max_rows + } else { + // Try the deprecated repr_rows attribute + let repr_rows = get_attr(formatter, "repr_rows", 0usize); + if repr_rows > 0 { + repr_rows + } else { + // Use default + default_config.max_rows + } + }; let config = FormatterConfig { max_bytes, min_rows, - repr_rows, + max_rows, }; // Return the validated config, converting String error to PyErr @@ -162,7 +190,13 @@ fn build_formatter_config_from_python(formatter: &Bound<'_, PyAny>) -> PyResult< } /// Python mapping of `ParquetOptions` (includes just the writer-related options). -#[pyclass(frozen, name = "ParquetWriterOptions", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ParquetWriterOptions", + module = "datafusion", + subclass +)] #[derive(Clone, Default)] pub struct PyParquetWriterOptions { options: ParquetOptions, @@ -175,7 +209,7 @@ impl PyParquetWriterOptions { pub fn new( data_pagesize_limit: usize, write_batch_size: usize, - writer_version: String, + writer_version: &str, skip_arrow_metadata: bool, compression: Option, dictionary_enabled: Option, @@ -193,8 +227,11 @@ impl PyParquetWriterOptions { allow_single_file_parallelism: bool, maximum_parallel_row_group_writers: usize, maximum_buffered_record_batches_per_stream: usize, - ) -> Self { - Self { + ) -> PyResult { + let writer_version = + datafusion::common::parquet_config::DFParquetWriterVersion::from_str(writer_version) + .map_err(py_datafusion_err)?; + Ok(Self { options: ParquetOptions { data_pagesize_limit, write_batch_size, @@ -218,12 +255,18 @@ impl PyParquetWriterOptions { maximum_buffered_record_batches_per_stream, ..Default::default() }, - } + }) } } /// Python mapping of `ParquetColumnOptions`. -#[pyclass(frozen, name = "ParquetColumnOptions", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ParquetColumnOptions", + module = "datafusion", + subclass +)] #[derive(Clone, Default)] pub struct PyParquetColumnOptions { options: ParquetColumnOptions, @@ -258,13 +301,22 @@ impl PyParquetColumnOptions { /// A PyDataFrame is a representation of a logical plan and an API to compose statements. /// Use it to build a plan and `.collect()` to execute the plan and collect the result. /// The actual execution of a plan runs natively on Rust and Arrow on a multi-threaded environment. -#[pyclass(name = "DataFrame", module = "datafusion", subclass, frozen)] +#[pyclass( + from_py_object, + name = "DataFrame", + module = "datafusion", + subclass, + frozen +)] #[derive(Clone)] pub struct PyDataFrame { df: Arc, // 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 { @@ -273,6 +325,7 @@ impl PyDataFrame { Self { df: Arc::new(df), batches: Arc::new(Mutex::new(None)), + last_plan: Arc::new(Mutex::new(None)), } } @@ -344,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 @@ -425,17 +492,17 @@ impl PyDataFrame { fn __getitem__(&self, key: Bound<'_, PyAny>) -> PyDataFusionResult { if let Ok(key) = key.extract::() { // df[col] - self.select_columns(vec![key]) - } else if let Ok(tuple) = key.downcast::() { + 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)) @@ -511,10 +578,8 @@ 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)?; + fn alias(&self, alias: &str) -> PyDataFusionResult { + let df = self.df.as_ref().clone().alias(alias)?; Ok(Self::new(df)) } @@ -539,6 +604,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)) @@ -602,8 +675,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() @@ -618,7 +692,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 @@ -637,7 +712,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) } @@ -761,9 +844,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) } @@ -778,7 +879,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()) } @@ -821,39 +928,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); - let df = self - .df - .as_ref() - .clone() - .unnest_columns_with_options(&[column], unnest_options)?; - 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 @@ -864,21 +946,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, @@ -1073,9 +1213,10 @@ impl PyDataFrame { let mut projection: Option = None; if let Some(schema_capsule) = requested_schema { - validate_pycapsule(&schema_capsule, "arrow_schema")?; - - let schema_ptr = unsafe { schema_capsule.reference::() }; + let data: NonNull = schema_capsule + .pointer_checked(Some(c"arrow_schema"))? + .cast(); + let schema_ptr = unsafe { data.as_ref() }; let desired_schema = Schema::try_from(schema_ptr)?; schema = project_schema(schema, desired_schema)?; @@ -1096,20 +1237,22 @@ impl PyDataFrame { // destructor provided by PyO3 will drop the stream unless ownership is // transferred to PyArrow during import. let stream = FFI_ArrowArrayStream::new(reader); - let name = CString::new(ARROW_ARRAY_STREAM_NAME.to_bytes()).unwrap(); - let capsule = PyCapsule::new(py, stream, Some(name))?; + let capsule = PyCapsule::new_with_value(py, stream, ARROW_ARRAY_STREAM_NAME)?; Ok(capsule) } 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()) } @@ -1166,20 +1309,28 @@ impl PyDataFrame { columns: Option>, py: Python, ) -> PyDataFusionResult { - let scalar_value = py_obj_to_scalar_value(py, value)?; + let scalar_value: PyScalarValue = value.extract(py)?; let cols = match columns { Some(col_names) => col_names.iter().map(|c| c.to_string()).collect(), None => Vec::new(), // Empty vector means fill null for all columns }; - let df = self.df.as_ref().clone().fill_null(scalar_value, cols)?; + let cols = cols.iter().map(String::as_str).collect::>(); + let df = self.df.as_ref().fill_null(&scalar_value.0, &cols)?; Ok(Self::new(df)) } } #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(frozen, eq, eq_int, name = "InsertOp", module = "datafusion")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "InsertOp", + module = "datafusion" +)] pub enum PyInsertOp { APPEND, REPLACE, @@ -1197,7 +1348,12 @@ impl From for InsertOp { } #[derive(Debug, Clone)] -#[pyclass(frozen, name = "DataFrameWriteOptions", module = "datafusion")] +#[pyclass( + from_py_object, + frozen, + name = "DataFrameWriteOptions", + module = "datafusion" +)] pub struct PyDataFrameWriteOptions { insert_operation: InsertOp, single_file_output: bool, @@ -1239,6 +1395,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 @@ -1303,7 +1479,10 @@ fn record_batch_into_schema( } else if field.is_nullable() { data_arrays.push(new_null_array(desired_data_type, array_size)); } else { - return Err(ArrowError::CastError(format!("Attempting to cast to non-nullable and non-castable field {} during schema projection.", field.name()))); + return Err(ArrowError::CastError(format!( + "Attempting to cast to non-nullable and non-castable field {} during schema projection.", + field.name() + ))); } } else { if !field.is_nullable() { @@ -1336,7 +1515,7 @@ async fn collect_record_batches_to_display( let FormatterConfig { max_bytes, min_rows, - repr_rows, + max_rows, } = config; let partitioned_stream = df.execute_stream_partitioned().await?; @@ -1346,8 +1525,11 @@ async fn collect_record_batches_to_display( let mut record_batches = Vec::default(); let mut has_more = false; - // ensure minimum rows even if memory/row limits are hit - while (size_estimate_so_far < max_bytes && rows_so_far < repr_rows) || rows_so_far < min_rows { + // Collect rows until we hit a limit (memory or max_rows) OR reach the guaranteed minimum. + // The minimum rows constraint overrides both memory and row limits to ensure a baseline + // of data is always displayed, even if it temporarily exceeds those limits. + // This provides better UX by guaranteeing users see at least min_rows rows. + while (size_estimate_so_far < max_bytes && rows_so_far < max_rows) || rows_so_far < min_rows { let mut rb = match stream.next().await { None => { break; @@ -1360,11 +1542,14 @@ async fn collect_record_batches_to_display( if rows_in_rb > 0 { size_estimate_so_far += rb.get_array_memory_size(); + // When memory limit is exceeded, scale back row count proportionally to stay within budget if size_estimate_so_far > max_bytes { let ratio = max_bytes as f32 / size_estimate_so_far as f32; let total_rows = rows_in_rb + rows_so_far; + // Calculate reduced rows maintaining the memory/data proportion let mut reduced_row_num = (total_rows as f32 * ratio).round() as usize; + // Ensure we always respect the minimum rows guarantee if reduced_row_num < min_rows { reduced_row_num = min_rows.min(total_rows); } @@ -1377,8 +1562,8 @@ async fn collect_record_batches_to_display( } } - if rows_in_rb + rows_so_far > repr_rows { - rb = rb.slice(0, repr_rows - rows_so_far); + if rows_in_rb + rows_so_far > max_rows { + rb = rb.slice(0, max_rows - rows_so_far); has_more = true; } diff --git a/src/dataset.rs b/crates/core/src/dataset.rs similarity index 94% rename from src/dataset.rs rename to crates/core/src/dataset.rs index 6a4fdb1fa..2a5770338 100644 --- a/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; @@ -47,7 +46,7 @@ impl Dataset { // Ensure that we were passed an instance of pyarrow.dataset.Dataset let ds = PyModule::import(py, "pyarrow.dataset")?; let ds_attr = ds.getattr("Dataset")?; - let ds_type = ds_attr.downcast::()?; + let ds_type = ds_attr.cast::()?; if dataset.is_instance(ds_type)? { Ok(Dataset { dataset: dataset.clone().unbind(), @@ -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/src/dataset_exec.rs b/crates/core/src/dataset_exec.rs similarity index 91% rename from src/dataset_exec.rs rename to crates/core/src/dataset_exec.rs index a83b10941..963603c8b 100644 --- a/src/dataset_exec.rs +++ b/crates/core/src/dataset_exec.rs @@ -15,25 +15,25 @@ // 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::utils::conjunction; use datafusion::logical_expr::Expr; -use datafusion::physical_expr::{EquivalenceProperties, LexOrdering}; +use datafusion::logical_expr::utils::conjunction; +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::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - SendableRecordBatchStream, Statistics, + PlanProperties, SendableRecordBatchStream, Statistics, }; -use futures::{stream, TryStreamExt}; +use futures::{TryStreamExt, stream}; /// Implements a Datafusion physical ExecutionPlan that delegates to a PyArrow Dataset /// This actually performs the projection, filtering and scanning of a Dataset use pyo3::prelude::*; @@ -71,7 +71,7 @@ pub(crate) struct DatasetExec { columns: Option>, filter_expr: Option>, projected_statistics: Statistics, - plan_properties: datafusion::physical_plan::PlanProperties, + plan_properties: Arc, } impl DatasetExec { @@ -111,7 +111,7 @@ impl DatasetExec { let scanner = dataset.call_method("scanner", (), Some(&kwargs))?; - let schema = Arc::new( + let schema: SchemaRef = Arc::new( scanner .getattr("projected_schema")? .extract::>()? @@ -128,15 +128,15 @@ impl DatasetExec { )?; let fragments_iter = pylist.call1((fragments_iterator,))?; - let fragments = fragments_iter.downcast::().map_err(PyErr::from)?; + let fragments = fragments_iter.cast::().map_err(PyErr::from)?; let projected_statistics = Statistics::new_unknown(&schema); - let plan_properties = datafusion::physical_plan::PlanProperties::new( + let plan_properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(schema.clone()), Partitioning::UnknownPartitioning(fragments.len()), EmissionType::Final, Boundedness::Bounded, - ); + )); Ok(DatasetExec { dataset: dataset.clone().unbind(), @@ -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() @@ -171,6 +166,15 @@ impl ExecutionPlan for DatasetExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DFResult, + ) -> DFResult { + // Any filters are pushed down into the PyArrow dataset scanner as pyarrow + // expressions, so this node owns no physical expressions to visit. + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, @@ -235,11 +239,11 @@ impl ExecutionPlan for DatasetExec { }) } - fn statistics(&self) -> DFResult { - Ok(self.projected_statistics.clone()) + fn partition_statistics(&self, _partition: Option) -> DFResult> { + Ok(Arc::new(self.projected_statistics.clone())) } - fn properties(&self) -> &datafusion::physical_plan::PlanProperties { + fn properties(&self) -> &Arc { &self.plan_properties } } diff --git a/src/expr/grouping_set.rs b/crates/core/src/errors.rs similarity index 61% rename from src/expr/grouping_set.rs rename to crates/core/src/errors.rs index 107dd9370..8babc5a56 100644 --- a/src/expr/grouping_set.rs +++ b/crates/core/src/errors.rs @@ -15,23 +15,4 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::GroupingSet; -use pyo3::prelude::*; - -#[pyclass(frozen, name = "GroupingSet", module = "datafusion.expr", subclass)] -#[derive(Clone)] -pub struct PyGroupingSet { - grouping_set: GroupingSet, -} - -impl From for GroupingSet { - fn from(grouping_set: PyGroupingSet) -> Self { - grouping_set.grouping_set - } -} - -impl From for PyGroupingSet { - fn from(grouping_set: GroupingSet) -> PyGroupingSet { - PyGroupingSet { grouping_set } - } -} +pub use datafusion_python_util::errors::*; diff --git a/src/expr.rs b/crates/core/src/expr.rs similarity index 82% rename from src/expr.rs rename to crates/core/src/expr.rs index fc8023b20..cab997d7a 100644 --- a/src/expr.rs +++ b/crates/core/src/expr.rs @@ -23,17 +23,21 @@ 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, - WindowFunction, + AggregateFunction, AggregateFunctionParams, FieldMetadata, HigherOrderFunction, InList, + InSubquery, Lambda, ScalarFunction, SetComparison, WindowFunction, }; use datafusion::logical_expr::utils::exprlist_to_fields; use datafusion::logical_expr::{ - col, lit, lit_with_metadata, Between, BinaryExpr, Case, Cast, Expr, ExprFuncBuilder, - ExprFunctionExt, Like, LogicalPlan, Operator, TryCast, WindowFunctionDefinition, + 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::IntoPyObjectExt; +use pyo3::types::PyBytes; use window::PyWindowFrame; use self::alias::PyAlias; @@ -43,8 +47,10 @@ 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::errors::{py_runtime_err, py_type_err, py_unsupported_variant_err, PyDataFusionResult}; +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; use crate::expr::column::PyColumn; @@ -85,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; @@ -98,6 +107,7 @@ pub mod recursive_query; pub mod repartition; pub mod scalar_subquery; pub mod scalar_variable; +pub mod set_comparison; pub mod signature; pub mod sort; pub mod sort_expr; @@ -111,10 +121,16 @@ pub mod unnest_expr; pub mod values; pub mod window; -use sort_expr::{to_sort_expressions, PySortExpr}; +use sort_expr::{PySortExpr, to_sort_expressions}; /// A PyExpr that can be used on a DataFrame -#[pyclass(frozen, name = "RawExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "RawExpr", + module = "datafusion.expr", + subclass +)] #[derive(Debug, Clone)] pub struct PyExpr { pub expr: Expr, @@ -141,15 +157,18 @@ pub fn py_expr_list(expr: &[Expr]) -> PyResult> { impl PyExpr { /// Return the specific expression fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { - Python::attach(|_| { - match &self.expr { + Python::attach(|_| match &self.expr { Expr::Alias(alias) => Ok(PyAlias::from(alias.clone()).into_bound_py_any(py)?), Expr::Column(col) => Ok(PyColumn::from(col.clone()).into_bound_py_any(py)?), - Expr::ScalarVariable(data_type, variables) => { - Ok(PyScalarVariable::new(data_type, variables).into_bound_py_any(py)?) + Expr::ScalarVariable(field, variables) => { + Ok(PyScalarVariable::new(field, variables).into_bound_py_any(py)?) } Expr::Like(value) => Ok(PyLike::from(value.clone()).into_bound_py_any(py)?), - Expr::Literal(value, metadata) => Ok(PyLiteral::new_with_metadata(value.clone(), metadata.clone()).into_bound_py_any(py)?), + Expr::Literal(value, metadata) => Ok(PyLiteral::new_with_metadata( + value.clone(), + metadata.clone(), + ) + .into_bound_py_any(py)?), Expr::BinaryExpr(expr) => Ok(PyBinaryExpr::from(expr.clone()).into_bound_py_any(py)?), Expr::Not(expr) => Ok(PyNot::new(*expr.clone()).into_bound_py_any(py)?), Expr::IsNotNull(expr) => Ok(PyIsNotNull::new(*expr.clone()).into_bound_py_any(py)?), @@ -159,13 +178,17 @@ impl PyExpr { Expr::IsUnknown(expr) => Ok(PyIsUnknown::new(*expr.clone()).into_bound_py_any(py)?), Expr::IsNotTrue(expr) => Ok(PyIsNotTrue::new(*expr.clone()).into_bound_py_any(py)?), Expr::IsNotFalse(expr) => Ok(PyIsNotFalse::new(*expr.clone()).into_bound_py_any(py)?), - Expr::IsNotUnknown(expr) => Ok(PyIsNotUnknown::new(*expr.clone()).into_bound_py_any(py)?), + Expr::IsNotUnknown(expr) => { + Ok(PyIsNotUnknown::new(*expr.clone()).into_bound_py_any(py)?) + } Expr::Negative(expr) => Ok(PyNegative::new(*expr.clone()).into_bound_py_any(py)?), Expr::AggregateFunction(expr) => { Ok(PyAggregateFunction::from(expr.clone()).into_bound_py_any(py)?) } Expr::SimilarTo(value) => Ok(PySimilarTo::from(value.clone()).into_bound_py_any(py)?), - Expr::Between(value) => Ok(between::PyBetween::from(value.clone()).into_bound_py_any(py)?), + Expr::Between(value) => { + Ok(between::PyBetween::from(value.clone()).into_bound_py_any(py)?) + } Expr::Case(value) => Ok(case::PyCase::from(value.clone()).into_bound_py_any(py)?), Expr::Cast(value) => Ok(cast::PyCast::from(value.clone()).into_bound_py_any(py)?), Expr::TryCast(value) => Ok(cast::PyTryCast::from(value.clone()).into_bound_py_any(py)?), @@ -175,7 +198,9 @@ impl PyExpr { Expr::WindowFunction(value) => Err(py_unsupported_variant_err(format!( "Converting Expr::WindowFunction to a Python object is not implemented: {value:?}" ))), - Expr::InList(value) => Ok(in_list::PyInList::from(value.clone()).into_bound_py_any(py)?), + Expr::InList(value) => { + Ok(in_list::PyInList::from(value.clone()).into_bound_py_any(py)?) + } Expr::Exists(value) => Ok(exists::PyExists::from(value.clone()).into_bound_py_any(py)?), Expr::InSubquery(value) => { Ok(in_subquery::PyInSubquery::from(value.clone()).into_bound_py_any(py)?) @@ -193,11 +218,25 @@ impl PyExpr { Expr::Placeholder(value) => { Ok(placeholder::PyPlaceholder::from(value.clone()).into_bound_py_any(py)?) } - Expr::OuterReferenceColumn(data_type, column) => Err(py_unsupported_variant_err(format!( - "Converting Expr::OuterReferenceColumn to a Python object is not implemented: {data_type:?} - {column:?}" - ))), - Expr::Unnest(value) => Ok(unnest_expr::PyUnnestExpr::from(value.clone()).into_bound_py_any(py)?), - } + Expr::OuterReferenceColumn(data_type, column) => { + Err(py_unsupported_variant_err(format!( + "Converting Expr::OuterReferenceColumn to a Python object is not implemented: {data_type:?} - {column:?}" + ))) + } + Expr::Unnest(value) => { + Ok(unnest_expr::PyUnnestExpr::from(value.clone()).into_bound_py_any(py)?) + } + 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)?) + } }) } @@ -319,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( @@ -364,11 +408,15 @@ impl PyExpr { | Expr::Placeholder { .. } | Expr::OuterReferenceColumn(_, _) | Expr::Unnest(_) - | Expr::IsNotUnknown(_) => RexType::Call, + | Expr::IsNotUnknown(_) + | Expr::SetComparison(_) + | Expr::HigherOrderFunction(..) + | Expr::Lambda(..) => RexType::Call, + Expr::LambdaVariable(..) => RexType::Reference, Expr::ScalarSubquery(..) => RexType::ScalarSubquery, #[allow(deprecated)] Expr::Wildcard { .. } => { - return Err(py_unsupported_variant_err("Expr::Wildcard is unsupported")) + return Err(py_unsupported_variant_err("Expr::Wildcard is unsupported")); } }) } @@ -385,7 +433,7 @@ impl PyExpr { Expr::Literal(scalar_value, _) => scalar_to_pyarrow(scalar_value, py), _ => Err(py_type_err(format!( "Non Expr::Literal encountered in types: {:?}", - &self.expr + self.expr ))), } } @@ -396,9 +444,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())]), @@ -415,20 +464,25 @@ impl PyExpr { | Expr::Negative(expr) | Expr::Cast(Cast { expr, .. }) | Expr::TryCast(TryCast { expr, .. }) - | Expr::InSubquery(InSubquery { expr, .. }) => Ok(vec![PyExpr::from(*expr.clone())]), + | Expr::InSubquery(InSubquery { expr, .. }) + | Expr::SetComparison(SetComparison { expr, .. }) => { + Ok(vec![PyExpr::from(*expr.clone())]) + } // Expr variants containing a collection of Expr(s) for operands Expr::AggregateFunction(AggregateFunction { 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 { @@ -518,6 +572,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(), @@ -554,8 +612,8 @@ impl PyExpr { _ => { return Err(py_type_err(format!( "Catch all triggered in get_operator_name: {:?}", - &self.expr - ))) + self.expr + ))); } }) } @@ -634,9 +692,64 @@ 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(frozen, name = "ExprFuncBuilder", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ExprFuncBuilder", + module = "datafusion.expr", + subclass +)] #[derive(Debug, Clone)] pub struct PyExprFuncBuilder { pub builder: ExprFuncBuilder, @@ -747,9 +860,12 @@ impl PyExpr { | Operator::AtQuestion | Operator::Question | Operator::QuestionAnd - | Operator::QuestionPipe => Err(py_type_err(format!("Unsupported expr: ${op}"))), + | 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:?}" @@ -805,6 +921,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/src/expr/aggregate.rs b/crates/core/src/expr/aggregate.rs similarity index 96% rename from src/expr/aggregate.rs rename to crates/core/src/expr/aggregate.rs index 4cb41b26a..7177fb469 100644 --- a/src/expr/aggregate.rs +++ b/crates/core/src/expr/aggregate.rs @@ -18,11 +18,11 @@ use std::fmt::{self, Display, Formatter}; use datafusion::common::DataFusionError; +use datafusion::logical_expr::Expr; use datafusion::logical_expr::expr::{AggregateFunction, AggregateFunctionParams, Alias}; use datafusion::logical_expr::logical_plan::Aggregate; -use datafusion::logical_expr::Expr; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; @@ -30,7 +30,13 @@ use crate::errors::py_type_err; use crate::expr::PyExpr; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Aggregate", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Aggregate", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyAggregate { aggregate: Aggregate, @@ -59,8 +65,8 @@ impl Display for PyAggregate { \nAggregates(s): {:?} \nInput: {:?} \nProjected Schema: {:?}", - &self.aggregate.group_expr, - &self.aggregate.aggr_expr, + self.aggregate.group_expr, + self.aggregate.aggr_expr, self.aggregate.input, self.aggregate.schema ) diff --git a/src/expr/aggregate_expr.rs b/crates/core/src/expr/aggregate_expr.rs similarity index 99% rename from src/expr/aggregate_expr.rs rename to crates/core/src/expr/aggregate_expr.rs index d3b695a27..88e47999f 100644 --- a/src/expr/aggregate_expr.rs +++ b/crates/core/src/expr/aggregate_expr.rs @@ -23,6 +23,7 @@ use pyo3::prelude::*; use crate::expr::PyExpr; #[pyclass( + from_py_object, frozen, name = "AggregateFunction", module = "datafusion.expr", diff --git a/src/expr/alias.rs b/crates/core/src/expr/alias.rs similarity index 91% rename from src/expr/alias.rs rename to crates/core/src/expr/alias.rs index c6d486284..391c94dbc 100644 --- a/src/expr/alias.rs +++ b/crates/core/src/expr/alias.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use crate::expr::PyExpr; -#[pyclass(frozen, name = "Alias", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Alias", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyAlias { alias: Alias, @@ -47,7 +53,7 @@ impl Display for PyAlias { "Alias \nExpr: `{:?}` \nAlias Name: `{}`", - &self.alias.expr, &self.alias.name + self.alias.expr, self.alias.name ) } } diff --git a/src/expr/analyze.rs b/crates/core/src/expr/analyze.rs similarity index 95% rename from src/expr/analyze.rs rename to crates/core/src/expr/analyze.rs index 05ec8dc22..137765fe1 100644 --- a/src/expr/analyze.rs +++ b/crates/core/src/expr/analyze.rs @@ -18,14 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::Analyze; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Analyze", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Analyze", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyAnalyze { analyze: Analyze, diff --git a/src/expr/between.rs b/crates/core/src/expr/between.rs similarity index 90% rename from src/expr/between.rs rename to crates/core/src/expr/between.rs index 4f0b34add..80f6c70da 100644 --- a/src/expr/between.rs +++ b/crates/core/src/expr/between.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use crate::expr::PyExpr; -#[pyclass(frozen, name = "Between", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Between", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyBetween { between: Between, @@ -49,7 +55,7 @@ impl Display for PyBetween { Negated: {:?} Low: {:?} High: {:?}", - &self.between.expr, &self.between.negated, &self.between.low, &self.between.high + self.between.expr, self.between.negated, self.between.low, self.between.high ) } } diff --git a/src/expr/binary_expr.rs b/crates/core/src/expr/binary_expr.rs similarity index 93% rename from src/expr/binary_expr.rs rename to crates/core/src/expr/binary_expr.rs index f67a08c7c..2326ba705 100644 --- a/src/expr/binary_expr.rs +++ b/crates/core/src/expr/binary_expr.rs @@ -20,7 +20,13 @@ use pyo3::prelude::*; use crate::expr::PyExpr; -#[pyclass(frozen, name = "BinaryExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "BinaryExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyBinaryExpr { expr: BinaryExpr, diff --git a/src/expr/bool_expr.rs b/crates/core/src/expr/bool_expr.rs similarity index 79% rename from src/expr/bool_expr.rs rename to crates/core/src/expr/bool_expr.rs index abd259409..d1cd7bcf7 100644 --- a/src/expr/bool_expr.rs +++ b/crates/core/src/expr/bool_expr.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use super::PyExpr; -#[pyclass(frozen, name = "Not", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Not", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyNot { expr: Expr, @@ -40,7 +46,7 @@ impl Display for PyNot { f, "Not Expr: {}", - &self.expr + self.expr ) } } @@ -52,7 +58,13 @@ impl PyNot { } } -#[pyclass(frozen, name = "IsNotNull", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotNull", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotNull { expr: Expr, @@ -70,7 +82,7 @@ impl Display for PyIsNotNull { f, "IsNotNull Expr: {}", - &self.expr + self.expr ) } } @@ -82,7 +94,13 @@ impl PyIsNotNull { } } -#[pyclass(frozen, name = "IsNull", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNull", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNull { expr: Expr, @@ -100,7 +118,7 @@ impl Display for PyIsNull { f, "IsNull Expr: {}", - &self.expr + self.expr ) } } @@ -112,7 +130,13 @@ impl PyIsNull { } } -#[pyclass(frozen, name = "IsTrue", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsTrue", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsTrue { expr: Expr, @@ -130,7 +154,7 @@ impl Display for PyIsTrue { f, "IsTrue Expr: {}", - &self.expr + self.expr ) } } @@ -142,7 +166,13 @@ impl PyIsTrue { } } -#[pyclass(frozen, name = "IsFalse", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsFalse", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsFalse { expr: Expr, @@ -160,7 +190,7 @@ impl Display for PyIsFalse { f, "IsFalse Expr: {}", - &self.expr + self.expr ) } } @@ -172,7 +202,13 @@ impl PyIsFalse { } } -#[pyclass(frozen, name = "IsUnknown", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsUnknown", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsUnknown { expr: Expr, @@ -190,7 +226,7 @@ impl Display for PyIsUnknown { f, "IsUnknown Expr: {}", - &self.expr + self.expr ) } } @@ -202,7 +238,13 @@ impl PyIsUnknown { } } -#[pyclass(frozen, name = "IsNotTrue", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotTrue", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotTrue { expr: Expr, @@ -220,7 +262,7 @@ impl Display for PyIsNotTrue { f, "IsNotTrue Expr: {}", - &self.expr + self.expr ) } } @@ -232,7 +274,13 @@ impl PyIsNotTrue { } } -#[pyclass(frozen, name = "IsNotFalse", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotFalse", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotFalse { expr: Expr, @@ -250,7 +298,7 @@ impl Display for PyIsNotFalse { f, "IsNotFalse Expr: {}", - &self.expr + self.expr ) } } @@ -262,7 +310,13 @@ impl PyIsNotFalse { } } -#[pyclass(frozen, name = "IsNotUnknown", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotUnknown", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotUnknown { expr: Expr, @@ -280,7 +334,7 @@ impl Display for PyIsNotUnknown { f, "IsNotUnknown Expr: {}", - &self.expr + self.expr ) } } @@ -292,7 +346,13 @@ impl PyIsNotUnknown { } } -#[pyclass(frozen, name = "Negative", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Negative", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyNegative { expr: Expr, @@ -310,7 +370,7 @@ impl Display for PyNegative { f, "Negative Expr: {}", - &self.expr + self.expr ) } } diff --git a/src/expr/case.rs b/crates/core/src/expr/case.rs similarity index 93% rename from src/expr/case.rs rename to crates/core/src/expr/case.rs index b49c19081..4f00449d8 100644 --- a/src/expr/case.rs +++ b/crates/core/src/expr/case.rs @@ -20,7 +20,13 @@ use pyo3::prelude::*; use crate::expr::PyExpr; -#[pyclass(frozen, name = "Case", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Case", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCase { case: Case, diff --git a/src/expr/cast.rs b/crates/core/src/expr/cast.rs similarity index 85% rename from src/expr/cast.rs rename to crates/core/src/expr/cast.rs index 1aca9ea95..484d0c059 100644 --- a/src/expr/cast.rs +++ b/crates/core/src/expr/cast.rs @@ -21,7 +21,13 @@ use pyo3::prelude::*; use crate::common::data_type::PyDataType; use crate::expr::PyExpr; -#[pyclass(frozen, name = "Cast", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Cast", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCast { cast: Cast, @@ -46,11 +52,11 @@ impl PyCast { } fn data_type(&self) -> PyResult { - Ok(self.cast.data_type.clone().into()) + Ok(self.cast.field.data_type().clone().into()) } } -#[pyclass(name = "TryCast", module = "datafusion.expr", subclass)] +#[pyclass(from_py_object, name = "TryCast", module = "datafusion.expr", subclass)] #[derive(Clone)] pub struct PyTryCast { try_cast: TryCast, @@ -75,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/src/expr/column.rs b/crates/core/src/expr/column.rs similarity index 93% rename from src/expr/column.rs rename to crates/core/src/expr/column.rs index 300079481..c1238f98a 100644 --- a/src/expr/column.rs +++ b/crates/core/src/expr/column.rs @@ -18,7 +18,13 @@ use datafusion::common::Column; use pyo3::prelude::*; -#[pyclass(frozen, name = "Column", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Column", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyColumn { pub col: Column, diff --git a/src/expr/conditional_expr.rs b/crates/core/src/expr/conditional_expr.rs similarity index 95% rename from src/expr/conditional_expr.rs rename to crates/core/src/expr/conditional_expr.rs index da6102dbf..ea21fdb20 100644 --- a/src/expr/conditional_expr.rs +++ b/crates/core/src/expr/conditional_expr.rs @@ -24,7 +24,13 @@ use crate::expr::PyExpr; // TODO(tsaucer) replace this all with CaseBuilder after it implements Clone #[derive(Clone, Debug)] -#[pyclass(name = "CaseBuilder", module = "datafusion.expr", subclass, frozen)] +#[pyclass( + from_py_object, + name = "CaseBuilder", + module = "datafusion.expr", + subclass, + frozen +)] pub struct PyCaseBuilder { expr: Option, when: Vec, diff --git a/src/expr/copy_to.rs b/crates/core/src/expr/copy_to.rs similarity index 93% rename from src/expr/copy_to.rs rename to crates/core/src/expr/copy_to.rs index 0b874e37d..78e53cdff 100644 --- a/src/expr/copy_to.rs +++ b/crates/core/src/expr/copy_to.rs @@ -21,13 +21,19 @@ use std::sync::Arc; use datafusion::common::file_options::file_type::FileType; use datafusion::logical_expr::dml::CopyTo; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "CopyTo", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "CopyTo", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCopyTo { copy: CopyTo, @@ -113,7 +119,13 @@ impl PyCopyTo { } } -#[pyclass(frozen, name = "FileType", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "FileType", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyFileType { file_type: Arc, diff --git a/src/expr/create_catalog.rs b/crates/core/src/expr/create_catalog.rs similarity index 95% rename from src/expr/create_catalog.rs rename to crates/core/src/expr/create_catalog.rs index 400246a82..fa95980c0 100644 --- a/src/expr/create_catalog.rs +++ b/crates/core/src/expr/create_catalog.rs @@ -19,14 +19,20 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use datafusion::logical_expr::CreateCatalog; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "CreateCatalog", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "CreateCatalog", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCreateCatalog { create: CreateCatalog, diff --git a/src/expr/create_catalog_schema.rs b/crates/core/src/expr/create_catalog_schema.rs similarity index 99% rename from src/expr/create_catalog_schema.rs rename to crates/core/src/expr/create_catalog_schema.rs index 641e2116d..d836284a0 100644 --- a/src/expr/create_catalog_schema.rs +++ b/crates/core/src/expr/create_catalog_schema.rs @@ -19,14 +19,15 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use datafusion::logical_expr::CreateCatalogSchema; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; #[pyclass( + from_py_object, frozen, name = "CreateCatalogSchema", module = "datafusion.expr", diff --git a/src/expr/create_external_table.rs b/crates/core/src/expr/create_external_table.rs similarity index 97% rename from src/expr/create_external_table.rs rename to crates/core/src/expr/create_external_table.rs index 05f9249b0..e78b836bf 100644 --- a/src/expr/create_external_table.rs +++ b/crates/core/src/expr/create_external_table.rs @@ -20,8 +20,8 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use datafusion::logical_expr::CreateExternalTable; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use super::sort_expr::PySortExpr; @@ -31,6 +31,7 @@ use crate::expr::PyExpr; use crate::sql::logical::PyLogicalPlan; #[pyclass( + from_py_object, frozen, name = "CreateExternalTable", module = "datafusion.expr", @@ -87,7 +88,7 @@ impl PyCreateExternalTable { let create = CreateExternalTable { schema: Arc::new(schema.into()), name: name.into(), - location, + locations: vec![location], file_type, table_partition_cols, if_not_exists, @@ -117,8 +118,8 @@ impl PyCreateExternalTable { Ok(self.create.name.to_string()) } - pub fn location(&self) -> String { - self.create.location.clone() + pub fn locations(&self) -> Vec { + self.create.locations.clone() } pub fn file_type(&self) -> String { diff --git a/src/expr/create_function.rs b/crates/core/src/expr/create_function.rs similarity index 94% rename from src/expr/create_function.rs rename to crates/core/src/expr/create_function.rs index 2a35635c2..622858913 100644 --- a/src/expr/create_function.rs +++ b/crates/core/src/expr/create_function.rs @@ -21,16 +21,22 @@ use std::sync::Arc; use datafusion::logical_expr::{ CreateFunction, CreateFunctionBody, OperateFunctionArg, Volatility, }; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; -use super::logical_node::LogicalNode; use super::PyExpr; +use super::logical_node::LogicalNode; use crate::common::data_type::PyDataType; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "CreateFunction", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "CreateFunction", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCreateFunction { create: CreateFunction, @@ -55,6 +61,7 @@ impl Display for PyCreateFunction { } #[pyclass( + from_py_object, frozen, name = "OperateFunctionArg", module = "datafusion.expr", @@ -66,7 +73,14 @@ pub struct PyOperateFunctionArg { } #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(frozen, eq, eq_int, name = "Volatility", module = "datafusion.expr")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "Volatility", + module = "datafusion.expr" +)] pub enum PyVolatility { Immutable, Stable, @@ -74,6 +88,7 @@ pub enum PyVolatility { } #[pyclass( + from_py_object, frozen, name = "CreateFunctionBody", module = "datafusion.expr", diff --git a/src/expr/create_index.rs b/crates/core/src/expr/create_index.rs similarity index 96% rename from src/expr/create_index.rs rename to crates/core/src/expr/create_index.rs index 5c378332c..5f9bd11e8 100644 --- a/src/expr/create_index.rs +++ b/crates/core/src/expr/create_index.rs @@ -19,15 +19,21 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use datafusion::logical_expr::CreateIndex; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use super::sort_expr::PySortExpr; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "CreateIndex", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "CreateIndex", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCreateIndex { create: CreateIndex, diff --git a/src/expr/create_memory_table.rs b/crates/core/src/expr/create_memory_table.rs similarity index 95% rename from src/expr/create_memory_table.rs rename to crates/core/src/expr/create_memory_table.rs index 7759eb420..c27c11834 100644 --- a/src/expr/create_memory_table.rs +++ b/crates/core/src/expr/create_memory_table.rs @@ -18,13 +18,14 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::CreateMemoryTable; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; #[pyclass( + from_py_object, frozen, name = "CreateMemoryTable", module = "datafusion.expr", @@ -56,10 +57,7 @@ impl Display for PyCreateMemoryTable { Input: {:?} if_not_exists: {:?} or_replace: {:?}", - &self.create.name, - &self.create.input, - &self.create.if_not_exists, - &self.create.or_replace, + self.create.name, self.create.input, self.create.if_not_exists, self.create.or_replace, ) } } diff --git a/src/expr/create_view.rs b/crates/core/src/expr/create_view.rs similarity index 93% rename from src/expr/create_view.rs rename to crates/core/src/expr/create_view.rs index 16faaf9d5..2f0315c40 100644 --- a/src/expr/create_view.rs +++ b/crates/core/src/expr/create_view.rs @@ -18,14 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::{CreateView, DdlStatement, LogicalPlan}; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::errors::py_type_err; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "CreateView", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "CreateView", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCreateView { create: CreateView, @@ -52,7 +58,7 @@ impl Display for PyCreateView { input: {:?} or_replace: {:?} definition: {:?}", - &self.create.name, &self.create.input, &self.create.or_replace, &self.create.definition, + self.create.name, self.create.input, self.create.or_replace, self.create.definition, ) } } diff --git a/src/expr/describe_table.rs b/crates/core/src/expr/describe_table.rs similarity index 95% rename from src/expr/describe_table.rs rename to crates/core/src/expr/describe_table.rs index 9b139ed3b..73955bb34 100644 --- a/src/expr/describe_table.rs +++ b/crates/core/src/expr/describe_table.rs @@ -21,14 +21,20 @@ use std::sync::Arc; use arrow::datatypes::Schema; use arrow::pyarrow::PyArrowType; use datafusion::logical_expr::DescribeTable; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "DescribeTable", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DescribeTable", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDescribeTable { describe: DescribeTable, diff --git a/src/expr/distinct.rs b/crates/core/src/expr/distinct.rs similarity index 95% rename from src/expr/distinct.rs rename to crates/core/src/expr/distinct.rs index 1505ec3e6..68c2a17fe 100644 --- a/src/expr/distinct.rs +++ b/crates/core/src/expr/distinct.rs @@ -18,13 +18,19 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::Distinct; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Distinct", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Distinct", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDistinct { distinct: Distinct, diff --git a/src/expr/dml.rs b/crates/core/src/expr/dml.rs similarity index 73% rename from src/expr/dml.rs rename to crates/core/src/expr/dml.rs index 091dcbc18..5967d181e 100644 --- a/src/expr/dml.rs +++ b/crates/core/src/expr/dml.rs @@ -17,15 +17,22 @@ use datafusion::logical_expr::dml::InsertOp; use datafusion::logical_expr::{DmlStatement, WriteOp}; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyNotImplementedError; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::common::schema::PyTableSource; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "DmlStatement", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DmlStatement", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDmlStatement { dml: DmlStatement, @@ -65,8 +72,8 @@ impl PyDmlStatement { }) } - pub fn op(&self) -> PyWriteOp { - self.dml.op.clone().into() + pub fn op(&self) -> PyResult { + self.dml.op.clone().try_into() } pub fn input(&self) -> PyLogicalPlan { @@ -89,27 +96,38 @@ impl PyDmlStatement { } #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(eq, eq_int, name = "WriteOp", module = "datafusion.expr")] +#[pyclass( + from_py_object, + eq, + eq_int, + name = "WriteOp", + module = "datafusion.expr" +)] pub enum PyWriteOp { Append, Overwrite, Replace, - Update, Delete, Ctas, + Truncate, } -impl From for PyWriteOp { - fn from(write_op: WriteOp) -> Self { - match write_op { - WriteOp::Insert(InsertOp::Append) => PyWriteOp::Append, - WriteOp::Insert(InsertOp::Overwrite) => PyWriteOp::Overwrite, - WriteOp::Insert(InsertOp::Replace) => PyWriteOp::Replace, +impl TryFrom for PyWriteOp { + type Error = PyErr; - WriteOp::Update => PyWriteOp::Update, - WriteOp::Delete => PyWriteOp::Delete, - WriteOp::Ctas => PyWriteOp::Ctas, + fn try_from(write_op: WriteOp) -> Result { + match write_op { + WriteOp::Insert(InsertOp::Append) => Ok(PyWriteOp::Append), + WriteOp::Insert(InsertOp::Overwrite) => Ok(PyWriteOp::Overwrite), + WriteOp::Insert(InsertOp::Replace) => Ok(PyWriteOp::Replace), + WriteOp::Update => Ok(PyWriteOp::Update), + WriteOp::Delete => Ok(PyWriteOp::Delete), + WriteOp::Ctas => Ok(PyWriteOp::Ctas), + WriteOp::Truncate => Ok(PyWriteOp::Truncate), + unsupported => Err(PyNotImplementedError::new_err(format!( + "DataFusion write operation {unsupported:?} is not supported" + ))), } } } @@ -120,10 +138,10 @@ impl From for WriteOp { PyWriteOp::Append => WriteOp::Insert(InsertOp::Append), PyWriteOp::Overwrite => WriteOp::Insert(InsertOp::Overwrite), PyWriteOp::Replace => WriteOp::Insert(InsertOp::Replace), - PyWriteOp::Update => WriteOp::Update, PyWriteOp::Delete => WriteOp::Delete, PyWriteOp::Ctas => WriteOp::Ctas, + PyWriteOp::Truncate => WriteOp::Truncate, } } } diff --git a/src/expr/drop_catalog_schema.rs b/crates/core/src/expr/drop_catalog_schema.rs similarity index 97% rename from src/expr/drop_catalog_schema.rs rename to crates/core/src/expr/drop_catalog_schema.rs index db6041a1b..f349098b7 100644 --- a/src/expr/drop_catalog_schema.rs +++ b/crates/core/src/expr/drop_catalog_schema.rs @@ -18,18 +18,18 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; -use datafusion::common::SchemaReference; +use datafusion::common::{SchemaReference, TableReference}; use datafusion::logical_expr::DropCatalogSchema; -use datafusion::sql::TableReference; +use pyo3::IntoPyObjectExt; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::IntoPyObjectExt; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; #[pyclass( + from_py_object, frozen, name = "DropCatalogSchema", module = "datafusion.expr", diff --git a/src/expr/drop_function.rs b/crates/core/src/expr/drop_function.rs similarity index 95% rename from src/expr/drop_function.rs rename to crates/core/src/expr/drop_function.rs index 070d15783..0599dd49e 100644 --- a/src/expr/drop_function.rs +++ b/crates/core/src/expr/drop_function.rs @@ -19,14 +19,20 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use datafusion::logical_expr::DropFunction; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "DropFunction", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DropFunction", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDropFunction { drop: DropFunction, diff --git a/src/expr/drop_table.rs b/crates/core/src/expr/drop_table.rs similarity index 92% rename from src/expr/drop_table.rs rename to crates/core/src/expr/drop_table.rs index ffb56e4ed..156d3bcc2 100644 --- a/src/expr/drop_table.rs +++ b/crates/core/src/expr/drop_table.rs @@ -18,13 +18,19 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::DropTable; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "DropTable", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DropTable", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDropTable { drop: DropTable, @@ -50,7 +56,7 @@ impl Display for PyDropTable { name: {:?} if_exists: {:?} schema: {:?}", - &self.drop.name, &self.drop.if_exists, &self.drop.schema, + self.drop.name, self.drop.if_exists, self.drop.schema, ) } } diff --git a/src/expr/drop_view.rs b/crates/core/src/expr/drop_view.rs similarity index 95% rename from src/expr/drop_view.rs rename to crates/core/src/expr/drop_view.rs index 9d72f2077..0d0c51f13 100644 --- a/src/expr/drop_view.rs +++ b/crates/core/src/expr/drop_view.rs @@ -19,14 +19,20 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use datafusion::logical_expr::DropView; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "DropView", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DropView", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDropView { drop: DropView, diff --git a/src/expr/empty_relation.rs b/crates/core/src/expr/empty_relation.rs similarity index 93% rename from src/expr/empty_relation.rs rename to crates/core/src/expr/empty_relation.rs index 35c3fa79b..5af4aafeb 100644 --- a/src/expr/empty_relation.rs +++ b/crates/core/src/expr/empty_relation.rs @@ -18,14 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::EmptyRelation; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "EmptyRelation", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "EmptyRelation", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyEmptyRelation { empty: EmptyRelation, @@ -50,7 +56,7 @@ impl Display for PyEmptyRelation { "Empty Relation Produce One Row: {:?} Schema: {:?}", - &self.empty.produce_one_row, &self.empty.schema + self.empty.produce_one_row, self.empty.schema ) } } diff --git a/src/expr/exists.rs b/crates/core/src/expr/exists.rs similarity index 91% rename from src/expr/exists.rs rename to crates/core/src/expr/exists.rs index 392bfcb9e..d2e816127 100644 --- a/src/expr/exists.rs +++ b/crates/core/src/expr/exists.rs @@ -20,7 +20,13 @@ use pyo3::prelude::*; use super::subquery::PySubquery; -#[pyclass(frozen, name = "Exists", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Exists", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyExists { exists: Exists, diff --git a/src/expr/explain.rs b/crates/core/src/expr/explain.rs similarity index 90% rename from src/expr/explain.rs rename to crates/core/src/expr/explain.rs index c6884e98a..d6ba1c25c 100644 --- a/src/expr/explain.rs +++ b/crates/core/src/expr/explain.rs @@ -17,17 +17,23 @@ use std::fmt::{self, Display, Formatter}; -use datafusion::logical_expr::logical_plan::Explain; use datafusion::logical_expr::LogicalPlan; -use pyo3::prelude::*; +use datafusion::logical_expr::logical_plan::Explain; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::errors::py_type_err; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Explain", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Explain", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyExplain { explain: Explain, @@ -55,11 +61,11 @@ impl Display for PyExplain { stringified_plans: {:?} schema: {:?} logical_optimization_succeeded: {:?}", - &self.explain.verbose, - &self.explain.plan, - &self.explain.stringified_plans, - &self.explain.schema, - &self.explain.logical_optimization_succeeded + self.explain.verbose, + self.explain.plan, + self.explain.stringified_plans, + self.explain.schema, + self.explain.logical_optimization_succeeded ) } } diff --git a/src/expr/extension.rs b/crates/core/src/expr/extension.rs similarity index 92% rename from src/expr/extension.rs rename to crates/core/src/expr/extension.rs index b4c688bd0..a0b617565 100644 --- a/src/expr/extension.rs +++ b/crates/core/src/expr/extension.rs @@ -16,13 +16,19 @@ // under the License. use datafusion::logical_expr::Extension; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Extension", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Extension", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyExtension { pub node: Extension, diff --git a/src/expr/filter.rs b/crates/core/src/expr/filter.rs similarity index 93% rename from src/expr/filter.rs rename to crates/core/src/expr/filter.rs index 25a1e76b3..1fe5f2c7f 100644 --- a/src/expr/filter.rs +++ b/crates/core/src/expr/filter.rs @@ -18,15 +18,21 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::Filter; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; -use crate::expr::logical_node::LogicalNode; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Filter", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Filter", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyFilter { filter: Filter, @@ -51,7 +57,7 @@ impl Display for PyFilter { "Filter Predicate: {:?} Input: {:?}", - &self.filter.predicate, &self.filter.input + self.filter.predicate, self.filter.input ) } } diff --git a/crates/core/src/expr/grouping_set.rs b/crates/core/src/expr/grouping_set.rs new file mode 100644 index 000000000..11d8f4fcd --- /dev/null +++ b/crates/core/src/expr/grouping_set.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 datafusion::logical_expr::{Expr, GroupingSet}; +use pyo3::prelude::*; + +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "GroupingSet", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +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 + } +} + +impl From for PyGroupingSet { + fn from(grouping_set: GroupingSet) -> PyGroupingSet { + PyGroupingSet { grouping_set } + } +} 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..5ba64052a --- /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/src/expr/in_list.rs b/crates/core/src/expr/in_list.rs similarity index 92% rename from src/expr/in_list.rs rename to crates/core/src/expr/in_list.rs index 128c3f4c2..0612cc21e 100644 --- a/src/expr/in_list.rs +++ b/crates/core/src/expr/in_list.rs @@ -20,7 +20,13 @@ use pyo3::prelude::*; use crate::expr::PyExpr; -#[pyclass(frozen, name = "InList", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "InList", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyInList { in_list: InList, diff --git a/src/expr/in_subquery.rs b/crates/core/src/expr/in_subquery.rs similarity index 92% rename from src/expr/in_subquery.rs rename to crates/core/src/expr/in_subquery.rs index 5cff86c06..81a2c5794 100644 --- a/src/expr/in_subquery.rs +++ b/crates/core/src/expr/in_subquery.rs @@ -18,10 +18,16 @@ use datafusion::logical_expr::expr::InSubquery; use pyo3::prelude::*; -use super::subquery::PySubquery; use super::PyExpr; +use super::subquery::PySubquery; -#[pyclass(frozen, name = "InSubquery", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "InSubquery", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyInSubquery { in_subquery: InSubquery, diff --git a/src/expr/indexed_field.rs b/crates/core/src/expr/indexed_field.rs similarity index 94% rename from src/expr/indexed_field.rs rename to crates/core/src/expr/indexed_field.rs index 1dfa0ed2f..98a90d8d4 100644 --- a/src/expr/indexed_field.rs +++ b/crates/core/src/expr/indexed_field.rs @@ -15,14 +15,21 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; +use std::fmt::{Display, Formatter}; + use datafusion::logical_expr::expr::{GetFieldAccess, GetIndexedField}; use pyo3::prelude::*; -use std::fmt::{Display, Formatter}; use super::literal::PyLiteral; +use crate::expr::PyExpr; -#[pyclass(frozen, name = "GetIndexedField", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "GetIndexedField", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyGetIndexedField { indexed_field: GetIndexedField, diff --git a/src/expr/join.rs b/crates/core/src/expr/join.rs similarity index 90% rename from src/expr/join.rs rename to crates/core/src/expr/join.rs index 82cc2a607..634e13734 100644 --- a/src/expr/join.rs +++ b/crates/core/src/expr/join.rs @@ -19,16 +19,16 @@ use std::fmt::{self, Display, Formatter}; use datafusion::common::NullEquality; use datafusion::logical_expr::logical_plan::{Join, JoinConstraint, JoinType}; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; -use crate::expr::logical_node::LogicalNode; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; #[derive(Debug, Clone, PartialEq, Eq, Hash)] -#[pyclass(frozen, name = "JoinType", module = "datafusion.expr")] +#[pyclass(from_py_object, frozen, name = "JoinType", module = "datafusion.expr")] pub struct PyJoinType { join_type: JoinType, } @@ -63,7 +63,12 @@ impl Display for PyJoinType { } #[derive(Debug, Clone, Copy)] -#[pyclass(frozen, name = "JoinConstraint", module = "datafusion.expr")] +#[pyclass( + from_py_object, + frozen, + name = "JoinConstraint", + module = "datafusion.expr" +)] pub struct PyJoinConstraint { join_constraint: JoinConstraint, } @@ -90,7 +95,13 @@ impl PyJoinConstraint { } } -#[pyclass(frozen, name = "Join", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Join", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyJoin { join: Join, @@ -121,14 +132,14 @@ impl Display for PyJoin { JoinConstraint: {:?} Schema: {:?} NullEquality: {:?}", - &self.join.left, - &self.join.right, - &self.join.on, - &self.join.filter, - &self.join.join_type, - &self.join.join_constraint, - &self.join.schema, - &self.join.null_equality, + self.join.left, + self.join.right, + self.join.on, + self.join.filter, + self.join.join_type, + self.join.join_constraint, + self.join.schema, + self.join.null_equality, ) } } diff --git a/crates/core/src/expr/lambda.rs b/crates/core/src/expr/lambda.rs new file mode 100644 index 000000000..7190521d6 --- /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..7baf5f21c --- /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/src/expr/like.rs b/crates/core/src/expr/like.rs similarity index 85% rename from src/expr/like.rs rename to crates/core/src/expr/like.rs index 94860bd6c..900551a20 100644 --- a/src/expr/like.rs +++ b/crates/core/src/expr/like.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use crate::expr::PyExpr; -#[pyclass(frozen, name = "Like", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Like", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyLike { like: Like, @@ -49,10 +55,10 @@ impl Display for PyLike { Expr: {:?} Pattern: {:?} Escape_Char: {:?}", - &self.negated(), - &self.expr(), - &self.pattern(), - &self.escape_char() + self.negated(), + self.expr(), + self.pattern(), + self.escape_char() ) } } @@ -80,7 +86,13 @@ impl PyLike { } } -#[pyclass(frozen, name = "ILike", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ILike", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyILike { like: Like, @@ -107,10 +119,10 @@ impl Display for PyILike { Expr: {:?} Pattern: {:?} Escape_Char: {:?}", - &self.negated(), - &self.expr(), - &self.pattern(), - &self.escape_char() + self.negated(), + self.expr(), + self.pattern(), + self.escape_char() ) } } @@ -138,7 +150,13 @@ impl PyILike { } } -#[pyclass(frozen, name = "SimilarTo", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SimilarTo", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySimilarTo { like: Like, @@ -165,10 +183,10 @@ impl Display for PySimilarTo { Expr: {:?} Pattern: {:?} Escape_Char: {:?}", - &self.negated(), - &self.expr(), - &self.pattern(), - &self.escape_char() + self.negated(), + self.expr(), + self.pattern(), + self.escape_char() ) } } diff --git a/src/expr/limit.rs b/crates/core/src/expr/limit.rs similarity index 70% rename from src/expr/limit.rs rename to crates/core/src/expr/limit.rs index 9318eff97..f53737b5e 100644 --- a/src/expr/limit.rs +++ b/crates/core/src/expr/limit.rs @@ -18,14 +18,21 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::Limit; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; +use crate::expr::PyExpr; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Limit", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Limit", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyLimit { limit: Limit, @@ -51,26 +58,30 @@ impl Display for PyLimit { Skip: {:?} Fetch: {:?} Input: {:?}", - &self.limit.skip, &self.limit.fetch, &self.limit.input + self.limit.skip, self.limit.fetch, self.limit.input ) } } #[pymethods] impl PyLimit { - // NOTE: Upstream now has expressions for skip and fetch - // TODO: Do we still want to expose these? - // REF: https://github.com/apache/datafusion/pull/12836 - - // /// Retrieves the skip value for this `Limit` - // fn skip(&self) -> usize { - // self.limit.skip - // } + // Retrieves the skip expression for this `Limit`, if any. + // + // `LIMIT`/`OFFSET` were changed upstream to support arbitrary + // expressions (not just constants), see + // https://github.com/apache/datafusion/pull/13028. Callers that expect + // a simple literal (the common case, e.g. `OFFSET 5`) should evaluate + // the returned `PyExpr` via `Expr.python_value()`. + fn skip(&self) -> PyResult> { + Ok(self.limit.skip.as_deref().cloned().map(PyExpr::from)) + } - // /// Retrieves the fetch value for this `Limit` - // fn fetch(&self) -> Option { - // self.limit.fetch - // } + // Retrieves the fetch expression for this `Limit`, if any. + // + // See the note on `skip` above regarding expression-based limits. + fn fetch(&self) -> PyResult> { + Ok(self.limit.fetch.as_deref().cloned().map(PyExpr::from)) + } /// Retrieves the input `LogicalPlan` to this `Limit` node fn input(&self) -> PyResult> { diff --git a/src/expr/literal.rs b/crates/core/src/expr/literal.rs similarity index 97% rename from src/expr/literal.rs rename to crates/core/src/expr/literal.rs index 3e8e229f9..9db0f594b 100644 --- a/src/expr/literal.rs +++ b/crates/core/src/expr/literal.rs @@ -17,12 +17,18 @@ use datafusion::common::ScalarValue; use datafusion::logical_expr::expr::FieldMetadata; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::errors::PyDataFusionError; -#[pyclass(name = "Literal", module = "datafusion.expr", subclass, frozen)] +#[pyclass( + from_py_object, + name = "Literal", + module = "datafusion.expr", + subclass, + frozen +)] #[derive(Clone)] pub struct PyLiteral { pub value: ScalarValue, diff --git a/src/expr/logical_node.rs b/crates/core/src/expr/logical_node.rs similarity index 100% rename from src/expr/logical_node.rs rename to crates/core/src/expr/logical_node.rs diff --git a/src/expr/placeholder.rs b/crates/core/src/expr/placeholder.rs similarity index 93% rename from src/expr/placeholder.rs rename to crates/core/src/expr/placeholder.rs index f1e8694a9..6bd88321c 100644 --- a/src/expr/placeholder.rs +++ b/crates/core/src/expr/placeholder.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use crate::common::data_type::PyDataType; -#[pyclass(frozen, name = "Placeholder", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Placeholder", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyPlaceholder { placeholder: Placeholder, diff --git a/src/expr/projection.rs b/crates/core/src/expr/projection.rs similarity index 94% rename from src/expr/projection.rs rename to crates/core/src/expr/projection.rs index bd21418a2..7e22e1e7b 100644 --- a/src/expr/projection.rs +++ b/crates/core/src/expr/projection.rs @@ -17,17 +17,23 @@ use std::fmt::{self, Display, Formatter}; -use datafusion::logical_expr::logical_plan::Projection; use datafusion::logical_expr::Expr; -use pyo3::prelude::*; +use datafusion::logical_expr::logical_plan::Projection; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; -use crate::expr::logical_node::LogicalNode; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Projection", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Projection", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyProjection { pub projection: Projection, @@ -59,7 +65,7 @@ impl Display for PyProjection { \nExpr(s): {:?} \nInput: {:?} \nProjected Schema: {:?}", - &self.projection.expr, &self.projection.input, &self.projection.schema, + self.projection.expr, self.projection.input, self.projection.schema, ) } } diff --git a/src/expr/recursive_query.rs b/crates/core/src/expr/recursive_query.rs similarity index 88% rename from src/expr/recursive_query.rs rename to crates/core/src/expr/recursive_query.rs index 0e1171ea9..0b198a191 100644 --- a/src/expr/recursive_query.rs +++ b/crates/core/src/expr/recursive_query.rs @@ -18,13 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::RecursiveQuery; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; +use crate::errors::PyDataFusionResult; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "RecursiveQuery", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "RecursiveQuery", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyRecursiveQuery { query: RecursiveQuery, @@ -61,15 +68,15 @@ impl PyRecursiveQuery { static_term: PyLogicalPlan, recursive_term: PyLogicalPlan, is_distinct: bool, - ) -> Self { - Self { - query: RecursiveQuery { + ) -> PyDataFusionResult { + Ok(Self { + query: RecursiveQuery::try_new( name, - static_term: static_term.plan(), - recursive_term: recursive_term.plan(), + static_term.plan(), + recursive_term.plan(), is_distinct, - }, - } + )?, + }) } fn name(&self) -> PyResult { diff --git a/src/expr/repartition.rs b/crates/core/src/expr/repartition.rs similarity index 92% rename from src/expr/repartition.rs rename to crates/core/src/expr/repartition.rs index 0b3cc4b2b..cbc8a97bb 100644 --- a/src/expr/repartition.rs +++ b/crates/core/src/expr/repartition.rs @@ -19,21 +19,33 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::Repartition; use datafusion::logical_expr::{Expr, Partitioning}; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; -use super::logical_node::LogicalNode; use super::PyExpr; +use super::logical_node::LogicalNode; use crate::errors::py_type_err; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Repartition", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Repartition", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyRepartition { repartition: Repartition, } -#[pyclass(frozen, name = "Partitioning", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Partitioning", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyPartitioning { partitioning: Partitioning, @@ -70,7 +82,7 @@ impl Display for PyRepartition { "Repartition input: {:?} partitioning_scheme: {:?}", - &self.repartition.input, &self.repartition.partitioning_scheme, + self.repartition.input, self.repartition.partitioning_scheme, ) } } diff --git a/src/expr/scalar_subquery.rs b/crates/core/src/expr/scalar_subquery.rs similarity index 91% rename from src/expr/scalar_subquery.rs rename to crates/core/src/expr/scalar_subquery.rs index e58d66e19..c7852a4c4 100644 --- a/src/expr/scalar_subquery.rs +++ b/crates/core/src/expr/scalar_subquery.rs @@ -20,7 +20,13 @@ use pyo3::prelude::*; use super::subquery::PySubquery; -#[pyclass(frozen, name = "ScalarSubquery", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ScalarSubquery", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyScalarSubquery { subquery: Subquery, diff --git a/src/expr/scalar_variable.rs b/crates/core/src/expr/scalar_variable.rs similarity index 76% rename from src/expr/scalar_variable.rs rename to crates/core/src/expr/scalar_variable.rs index f3c128a4c..2d3bc4b76 100644 --- a/src/expr/scalar_variable.rs +++ b/crates/core/src/expr/scalar_variable.rs @@ -15,22 +15,28 @@ // specific language governing permissions and limitations // under the License. -use datafusion::arrow::datatypes::DataType; +use arrow::datatypes::FieldRef; use pyo3::prelude::*; use crate::common::data_type::PyDataType; -#[pyclass(frozen, name = "ScalarVariable", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ScalarVariable", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyScalarVariable { - data_type: DataType, + field: FieldRef, variables: Vec, } impl PyScalarVariable { - pub fn new(data_type: &DataType, variables: &[String]) -> Self { + pub fn new(field: &FieldRef, variables: &[String]) -> Self { Self { - data_type: data_type.to_owned(), + field: field.to_owned(), variables: variables.to_vec(), } } @@ -40,7 +46,7 @@ impl PyScalarVariable { impl PyScalarVariable { /// Get the data type fn data_type(&self) -> PyResult { - Ok(self.data_type.clone().into()) + Ok(self.field.data_type().clone().into()) } fn variables(&self) -> PyResult> { @@ -48,6 +54,6 @@ impl PyScalarVariable { } fn __repr__(&self) -> PyResult { - Ok(format!("{}{:?}", self.data_type, self.variables)) + Ok(format!("{}{:?}", self.field.data_type(), self.variables)) } } diff --git a/crates/core/src/expr/set_comparison.rs b/crates/core/src/expr/set_comparison.rs new file mode 100644 index 000000000..9f0c077e1 --- /dev/null +++ b/crates/core/src/expr/set_comparison.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. + +use datafusion::logical_expr::expr::SetComparison; +use pyo3::prelude::*; + +use super::subquery::PySubquery; +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "SetComparison", + module = "datafusion.set_comparison", + subclass +)] +#[derive(Clone)] +pub struct PySetComparison { + set_comparison: SetComparison, +} + +impl From for PySetComparison { + fn from(set_comparison: SetComparison) -> Self { + PySetComparison { set_comparison } + } +} + +#[pymethods] +impl PySetComparison { + fn expr(&self) -> PyExpr { + (*self.set_comparison.expr).clone().into() + } + + fn subquery(&self) -> PySubquery { + self.set_comparison.subquery.clone().into() + } + + fn op(&self) -> String { + format!("{}", self.set_comparison.op) + } + + fn quantifier(&self) -> String { + format!("{}", self.set_comparison.quantifier) + } +} diff --git a/src/expr/signature.rs b/crates/core/src/expr/signature.rs similarity index 91% rename from src/expr/signature.rs rename to crates/core/src/expr/signature.rs index e2c23dce9..35268e3a9 100644 --- a/src/expr/signature.rs +++ b/crates/core/src/expr/signature.rs @@ -19,7 +19,13 @@ use datafusion::logical_expr::{TypeSignature, Volatility}; use pyo3::prelude::*; #[allow(dead_code)] -#[pyclass(frozen, name = "Signature", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Signature", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySignature { type_signature: TypeSignature, diff --git a/src/expr/sort.rs b/crates/core/src/expr/sort.rs similarity index 95% rename from src/expr/sort.rs rename to crates/core/src/expr/sort.rs index 8914c8f93..1b1065011 100644 --- a/src/expr/sort.rs +++ b/crates/core/src/expr/sort.rs @@ -19,15 +19,21 @@ use std::fmt::{self, Display, Formatter}; use datafusion::common::DataFusionError; use datafusion::logical_expr::logical_plan::Sort; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; use crate::expr::logical_node::LogicalNode; use crate::expr::sort_expr::PySortExpr; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Sort", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Sort", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySort { sort: Sort, @@ -55,7 +61,7 @@ impl Display for PySort { \nExpr(s): {:?} \nInput: {:?} \nSchema: {:?}", - &self.sort.expr, + self.sort.expr, self.sort.input, self.sort.input.schema() ) diff --git a/src/expr/sort_expr.rs b/crates/core/src/expr/sort_expr.rs similarity index 92% rename from src/expr/sort_expr.rs rename to crates/core/src/expr/sort_expr.rs index 23c066156..93faffcec 100644 --- a/src/expr/sort_expr.rs +++ b/crates/core/src/expr/sort_expr.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use crate::expr::PyExpr; -#[pyclass(frozen, name = "SortExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SortExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySortExpr { pub(crate) sort: SortExpr, @@ -48,7 +54,7 @@ impl Display for PySortExpr { Expr: {:?} Asc: {:?} NullsFirst: {:?}", - &self.sort.expr, &self.sort.asc, &self.sort.nulls_first + self.sort.expr, self.sort.asc, self.sort.nulls_first ) } } diff --git a/src/expr/statement.rs b/crates/core/src/expr/statement.rs similarity index 86% rename from src/expr/statement.rs rename to crates/core/src/expr/statement.rs index 40666dd8b..5aa1e4e9c 100644 --- a/src/expr/statement.rs +++ b/crates/core/src/expr/statement.rs @@ -20,17 +20,18 @@ use std::sync::Arc; use arrow::datatypes::Field; use arrow::pyarrow::PyArrowType; use datafusion::logical_expr::{ - Deallocate, Execute, Prepare, SetVariable, TransactionAccessMode, TransactionConclusion, - TransactionEnd, TransactionIsolationLevel, TransactionStart, + Deallocate, Execute, Prepare, ResetVariable, SetVariable, TransactionAccessMode, + TransactionConclusion, TransactionEnd, TransactionIsolationLevel, TransactionStart, }; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; -use super::logical_node::LogicalNode; use super::PyExpr; +use super::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; #[pyclass( + from_py_object, frozen, name = "TransactionStart", module = "datafusion.expr", @@ -67,6 +68,7 @@ impl LogicalNode for PyTransactionStart { #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] #[pyclass( + from_py_object, frozen, eq, eq_int, @@ -100,6 +102,7 @@ impl TryFrom for TransactionAccessMode { #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] #[pyclass( + from_py_object, frozen, eq, eq_int, @@ -178,7 +181,13 @@ impl PyTransactionStart { } } -#[pyclass(frozen, name = "TransactionEnd", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "TransactionEnd", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyTransactionEnd { transaction_end: TransactionEnd, @@ -210,6 +219,7 @@ impl LogicalNode for PyTransactionEnd { #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] #[pyclass( + from_py_object, frozen, eq, eq_int, @@ -259,7 +269,63 @@ impl PyTransactionEnd { } } -#[pyclass(frozen, name = "SetVariable", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ResetVariable", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyResetVariable { + reset_variable: ResetVariable, +} + +impl From for PyResetVariable { + fn from(reset_variable: ResetVariable) -> PyResetVariable { + PyResetVariable { reset_variable } + } +} + +impl TryFrom for ResetVariable { + type Error = PyErr; + + fn try_from(py: PyResetVariable) -> Result { + Ok(py.reset_variable) + } +} + +impl LogicalNode for PyResetVariable { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyResetVariable { + #[new] + pub fn new(variable: String) -> Self { + PyResetVariable { + reset_variable: ResetVariable { variable }, + } + } + + pub fn variable(&self) -> String { + self.reset_variable.variable.clone() + } +} + +#[pyclass( + from_py_object, + frozen, + name = "SetVariable", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySetVariable { set_variable: SetVariable, @@ -307,7 +373,13 @@ impl PySetVariable { } } -#[pyclass(frozen, name = "Prepare", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Prepare", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyPrepare { prepare: Prepare, @@ -372,7 +444,13 @@ impl PyPrepare { } } -#[pyclass(frozen, name = "Execute", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Execute", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyExecute { execute: Execute, @@ -429,7 +507,13 @@ impl PyExecute { } } -#[pyclass(frozen, name = "Deallocate", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Deallocate", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDeallocate { deallocate: Deallocate, diff --git a/src/expr/subquery.rs b/crates/core/src/expr/subquery.rs similarity index 95% rename from src/expr/subquery.rs rename to crates/core/src/expr/subquery.rs index 94c2583ba..c6fa83db8 100644 --- a/src/expr/subquery.rs +++ b/crates/core/src/expr/subquery.rs @@ -18,13 +18,19 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::Subquery; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Subquery", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Subquery", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySubquery { subquery: Subquery, diff --git a/src/expr/subquery_alias.rs b/crates/core/src/expr/subquery_alias.rs similarity index 95% rename from src/expr/subquery_alias.rs rename to crates/core/src/expr/subquery_alias.rs index 9bf1c9c51..a6b09e842 100644 --- a/src/expr/subquery_alias.rs +++ b/crates/core/src/expr/subquery_alias.rs @@ -18,14 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::SubqueryAlias; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "SubqueryAlias", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SubqueryAlias", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySubqueryAlias { subquery_alias: SubqueryAlias, diff --git a/src/expr/table_scan.rs b/crates/core/src/expr/table_scan.rs similarity index 95% rename from src/expr/table_scan.rs rename to crates/core/src/expr/table_scan.rs index bbf225f4c..46ce551d2 100644 --- a/src/expr/table_scan.rs +++ b/crates/core/src/expr/table_scan.rs @@ -19,15 +19,21 @@ use std::fmt::{self, Display, Formatter}; use datafusion::common::TableReference; use datafusion::logical_expr::logical_plan::TableScan; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; -use crate::expr::logical_node::LogicalNode; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "TableScan", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "TableScan", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyTableScan { table_scan: TableScan, @@ -59,10 +65,10 @@ impl Display for PyTableScan { Projections: {:?} Projected Schema: {:?} Filters: {:?}", - &self.table_scan.table_name, - &self.py_projections(), - &self.py_schema(), - &self.py_filters(), + self.table_scan.table_name, + self.py_projections(), + self.py_schema(), + self.py_filters(), ) } } diff --git a/src/expr/union.rs b/crates/core/src/expr/union.rs similarity index 93% rename from src/expr/union.rs rename to crates/core/src/expr/union.rs index c74d170aa..bd5770e0a 100644 --- a/src/expr/union.rs +++ b/crates/core/src/expr/union.rs @@ -18,14 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::Union; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Union", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Union", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyUnion { union_: Union, @@ -50,7 +56,7 @@ impl Display for PyUnion { "Union Inputs: {:?} Schema: {:?}", - &self.union_.inputs, &self.union_.schema, + self.union_.inputs, self.union_.schema, ) } } diff --git a/src/expr/unnest.rs b/crates/core/src/expr/unnest.rs similarity index 93% rename from src/expr/unnest.rs rename to crates/core/src/expr/unnest.rs index 7e68c15f4..540667824 100644 --- a/src/expr/unnest.rs +++ b/crates/core/src/expr/unnest.rs @@ -18,14 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::Unnest; -use pyo3::prelude::*; use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Unnest", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Unnest", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyUnnest { unnest_: Unnest, @@ -50,7 +56,7 @@ impl Display for PyUnnest { "Unnest Inputs: {:?} Schema: {:?}", - &self.unnest_.input, &self.unnest_.schema, + self.unnest_.input, self.unnest_.schema, ) } } diff --git a/src/expr/unnest_expr.rs b/crates/core/src/expr/unnest_expr.rs similarity index 92% rename from src/expr/unnest_expr.rs rename to crates/core/src/expr/unnest_expr.rs index dc6c4cb50..549257b86 100644 --- a/src/expr/unnest_expr.rs +++ b/crates/core/src/expr/unnest_expr.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use super::PyExpr; -#[pyclass(frozen, name = "UnnestExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "UnnestExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyUnnestExpr { unnest: Unnest, @@ -46,7 +52,7 @@ impl Display for PyUnnestExpr { f, "Unnest Expr: {:?}", - &self.unnest.expr, + self.unnest.expr, ) } } diff --git a/src/expr/values.rs b/crates/core/src/expr/values.rs similarity index 93% rename from src/expr/values.rs rename to crates/core/src/expr/values.rs index 7ae7350fc..d40b0e7cf 100644 --- a/src/expr/values.rs +++ b/crates/core/src/expr/values.rs @@ -19,14 +19,20 @@ use std::sync::Arc; use datafusion::logical_expr::Values; use pyo3::prelude::*; -use pyo3::{pyclass, IntoPyObjectExt, PyErr, PyResult, Python}; +use pyo3::{IntoPyObjectExt, PyErr, PyResult, Python, pyclass}; -use super::logical_node::LogicalNode; use super::PyExpr; +use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Values", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Values", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyValues { values: Values, diff --git a/src/expr/window.rs b/crates/core/src/expr/window.rs similarity index 95% rename from src/expr/window.rs rename to crates/core/src/expr/window.rs index b93e813c4..e0050f671 100644 --- a/src/expr/window.rs +++ b/crates/core/src/expr/window.rs @@ -19,26 +19,38 @@ use std::fmt::{self, Display, Formatter}; use datafusion::common::{DataFusionError, ScalarValue}; use datafusion::logical_expr::{Expr, Window, WindowFrame, WindowFrameBound, WindowFrameUnits}; +use pyo3::IntoPyObjectExt; use pyo3::exceptions::PyNotImplementedError; use pyo3::prelude::*; -use pyo3::IntoPyObjectExt; use super::py_expr_list; use crate::common::data_type::PyScalarValue; use crate::common::df_schema::PyDFSchema; -use crate::errors::{py_type_err, PyDataFusionResult}; -use crate::expr::logical_node::LogicalNode; -use crate::expr::sort_expr::{py_sort_expr_list, PySortExpr}; +use crate::errors::{PyDataFusionResult, py_type_err}; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; +use crate::expr::sort_expr::{PySortExpr, py_sort_expr_list}; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "WindowExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "WindowExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyWindowExpr { window: Window, } -#[pyclass(frozen, name = "WindowFrame", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "WindowFrame", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyWindowFrame { window_frame: WindowFrame, @@ -57,6 +69,7 @@ impl From for PyWindowFrame { } #[pyclass( + from_py_object, frozen, name = "WindowFrameBound", module = "datafusion.expr", @@ -92,7 +105,7 @@ impl Display for PyWindowExpr { "Over\n Window Expr: {:?} Schema: {:?}", - &self.window.window_expr, &self.window.schema + self.window.window_expr, self.window.schema ) } } diff --git a/src/functions.rs b/crates/core/src/functions.rs similarity index 76% rename from src/functions.rs rename to crates/core/src/functions.rs index e67781ccd..e57c7702d 100644 --- a/src/functions.rs +++ b/crates/core/src/functions.rs @@ -18,26 +18,20 @@ 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::{lit, Expr, ExprFunctionExt, WindowFrame, WindowFunctionDefinition}; +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::{to_sort_expressions, PySortExpr}; +use crate::expr::sort_expr::{PySortExpr, to_sort_expressions}; use crate::expr::window::PyWindowFrame; -use crate::expr::PyExpr; -fn add_builder_fns_to_aggregate( +pub(crate) fn add_builder_fns_to_aggregate( agg_fn: Expr, distinct: Option, filter: Option, @@ -93,6 +87,57 @@ 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(); + 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 { @@ -114,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)] @@ -189,6 +272,29 @@ fn regexp_count( .into()) } +#[pyfunction] +#[pyo3(signature = (values, regex, start=None, n=None, flags=None, subexpr=None))] +/// Returns the position in a string where the specified occurrence of a regular expression is located +fn regexp_instr( + values: PyExpr, + regex: PyExpr, + start: Option, + n: Option, + flags: Option, + subexpr: Option, +) -> PyResult { + Ok(functions::expr_fn::regexp_instr( + values.into(), + regex.into(), + start.map(|x| x.expr).or(Some(lit(1))), + n.map(|x| x.expr).or(Some(lit(1))), + None, + flags.map(|x| x.expr).or(Some(lit(""))), + subexpr.map(|x| x.expr).or(Some(lit(0))), + ) + .into()) +} + /// Creates a new Sort Expr #[pyfunction] fn order_by(expr: PyExpr, asc: bool, nulls_first: bool) -> PyResult { @@ -232,126 +338,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 @@ -441,7 +427,11 @@ macro_rules! array_fn { expr_fn!(abs, num); expr_fn!(acos, num); expr_fn!(acosh, num); -expr_fn!(ascii, arg1, "Returns the numeric code of the first character of the argument. In UTF8 encoding, returns the Unicode code point of the character. In other multibyte encodings, the argument must be an ASCII character."); +expr_fn!( + ascii, + arg1, + "Returns the numeric code of the first character of the argument. In UTF8 encoding, returns the Unicode code point of the character. In other multibyte encodings, the argument must be an ASCII character." +); expr_fn!(asin, num); expr_fn!(asinh, num); expr_fn!(atan, num); @@ -452,7 +442,10 @@ expr_fn!( arg, "Returns number of bits in the string (8 times the octet_length)." ); -expr_fn_vec!(btrim, "Removes the longest string containing only characters in characters (a space by default) from the start and end of string."); +expr_fn_vec!( + btrim, + "Removes the longest string containing only characters in characters (a space by default) from the start and end of string." +); expr_fn!(cbrt, num); expr_fn!(ceil, num); expr_fn!( @@ -464,6 +457,13 @@ 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, + "Return true if search_str is found within string (case-sensitive)." +); expr_fn!(cos, num); expr_fn!(cosh, num); expr_fn!(cot, num); @@ -475,7 +475,11 @@ expr_fn!(exp, num); expr_fn!(factorial, num); expr_fn!(floor, num); expr_fn!(gcd, x y); -expr_fn!(initcap, string, "Converts the first letter of each word to upper case and the rest to lower case. Words are sequences of alphanumeric characters separated by non-alphanumeric characters."); +expr_fn!( + initcap, + string, + "Converts the first letter of each word to upper case and the rest to lower case. Words are sequences of alphanumeric characters separated by non-alphanumeric characters." +); expr_fn!(isnan, num); expr_fn!(iszero, num); expr_fn!(levenshtein, string1 string2); @@ -486,8 +490,14 @@ expr_fn!(log, base num); expr_fn!(log10, num); expr_fn!(log2, num); expr_fn!(lower, arg1, "Converts the string to all lower case"); -expr_fn_vec!(lpad, "Extends the string to length length by prepending the characters fill (a space by default). If the string is already longer than length then it is truncated (on the right)."); -expr_fn_vec!(ltrim, "Removes the longest string containing only characters in characters (a space by default) from the start of string."); +expr_fn_vec!( + lpad, + "Extends the string to length length by prepending the characters fill (a space by default). If the string is already longer than length then it is truncated (on the right)." +); +expr_fn_vec!( + ltrim, + "Removes the longest string containing only characters in characters (a space by default) from the start of string." +); expr_fn!( md5, input_arg, @@ -503,8 +513,17 @@ 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, args, "Returns number of bytes in the string. Since this version of the function accepts type character directly, it will not strip trailing spaces."); +expr_fn!( + octet_length, + args, + "Returns number of bytes in the string. Since this version of the function accepts type character directly, it will not strip trailing spaces." +); expr_fn_vec!(overlay); expr_fn!(pi); expr_fn!(power, base exponent); @@ -522,8 +541,14 @@ expr_fn!( ); expr_fn!(right, string n, "Returns last n characters in the string, or when n is negative, returns all but first |n| characters."); expr_fn_vec!(round); -expr_fn_vec!(rpad, "Extends the string to length length by appending the characters fill (a space by default). If the string is already longer than length then it is truncated."); -expr_fn_vec!(rtrim, "Removes the longest string containing only characters in characters (a space by default) from the end of string."); +expr_fn_vec!( + rpad, + "Extends the string to length length by appending the characters fill (a space by default). If the string is already longer than length then it is truncated." +); +expr_fn_vec!( + rtrim, + "Removes the longest string containing only characters in characters (a space by default) from the end of string." +); expr_fn!(sha224, input_arg1); expr_fn!(sha256, input_arg1); expr_fn!(sha384, input_arg1); @@ -551,6 +576,9 @@ expr_fn!( "Converts the number to its equivalent hexadecimal representation." ); expr_fn!(now); +expr_fn_vec!(to_date); +expr_fn_vec!(to_local_time); +expr_fn_vec!(to_time); expr_fn_vec!(to_timestamp); expr_fn_vec!(to_timestamp_millis); expr_fn_vec!(to_timestamp_nanos); @@ -563,9 +591,14 @@ 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."); -expr_fn_vec!(trim, "Removes the longest string containing only characters in characters (a space by default) from the start, end, or both ends (BOTH is the default) of string."); +expr_fn_vec!( + trim, + "Removes the longest string containing only characters in characters (a space by default) from the start, end, or both ends (BOTH is the default) of string." +); expr_fn_vec!(trunc); expr_fn!(upper, arg1, "Converts the string to all upper case."); expr_fn!(uuid); @@ -574,8 +607,34 @@ 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); +#[pyfunction] +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] +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); @@ -600,14 +659,28 @@ 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); 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); +// 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); @@ -639,9 +712,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))] @@ -679,6 +753,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] @@ -879,10 +966,17 @@ 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))?; 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))?; @@ -903,6 +997,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))?; @@ -917,6 +1012,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))?; @@ -924,13 +1020,15 @@ 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!(grouping))?; + 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))?; m.add_wrapped(wrap_pyfunction!(isnan))?; 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))?; @@ -948,6 +1046,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))?; @@ -958,6 +1057,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(radians))?; m.add_wrapped(wrap_pyfunction!(random))?; m.add_wrapped(wrap_pyfunction!(regexp_count))?; + m.add_wrapped(wrap_pyfunction!(regexp_instr))?; m.add_wrapped(wrap_pyfunction!(regexp_like))?; m.add_wrapped(wrap_pyfunction!(regexp_match))?; m.add_wrapped(wrap_pyfunction!(regexp_replace))?; @@ -991,6 +1091,10 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(tan))?; m.add_wrapped(wrap_pyfunction!(tanh))?; m.add_wrapped(wrap_pyfunction!(to_hex))?; + m.add_wrapped(wrap_pyfunction!(to_char))?; + m.add_wrapped(wrap_pyfunction!(to_date))?; + m.add_wrapped(wrap_pyfunction!(to_local_time))?; + m.add_wrapped(wrap_pyfunction!(to_time))?; m.add_wrapped(wrap_pyfunction!(to_timestamp))?; m.add_wrapped(wrap_pyfunction!(to_timestamp_millis))?; m.add_wrapped(wrap_pyfunction!(to_timestamp_nanos))?; @@ -1001,10 +1105,13 @@ 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))?; - 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))?; @@ -1027,12 +1134,23 @@ 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))?; 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))?; @@ -1059,9 +1177,24 @@ 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))?; + // 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/src/lib.rs b/crates/core/src/lib.rs similarity index 82% rename from src/lib.rs rename to crates/core/src/lib.rs index 9483a5252..7f0f9cb39 100644 --- a/src/lib.rs +++ b/crates/core/src/lib.rs @@ -16,9 +16,9 @@ // under the License. // Re-export Apache Arrow DataFusion dependencies -pub use datafusion; pub use datafusion::{ - common as datafusion_common, logical_expr as datafusion_expr, optimizer, sql as datafusion_sql, + self, common as datafusion_common, logical_expr as datafusion_expr, optimizer, + sql as datafusion_sql, }; #[cfg(feature = "substrait")] pub use datafusion_substrait; @@ -26,49 +26,47 @@ pub use datafusion_substrait; use mimalloc::MiMalloc; use pyo3::prelude::*; -#[allow(clippy::borrow_deref_ref)] +pub mod analyzer; 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)] pub mod dataframe; mod dataset; mod dataset_exec; pub mod errors; -#[allow(clippy::borrow_deref_ref)] pub mod expr; -#[allow(clippy::borrow_deref_ref)] mod functions; +pub mod metrics; +mod options; pub mod physical_plan; mod pyarrow_filter_expression; pub mod pyarrow_util; mod record_batch; +mod spark_functions; pub mod sql; pub mod store; pub mod table; pub mod unparser; +mod array; #[cfg(feature = "substrait")] pub mod substrait; -#[allow(clippy::borrow_deref_ref)] mod udaf; -#[allow(clippy::borrow_deref_ref)] mod udf; pub mod udtf; mod udwf; -pub mod utils; + +// 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; -// Used to define Tokio Runtime as a Python module attribute -pub(crate) struct TokioRuntime(tokio::runtime::Runtime); - /// Low-level DataFusion internal package. /// /// The higher-level public API is defined in pure python files under the @@ -92,8 +90,9 @@ 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::()?; m.add_class::()?; m.add_class::()?; @@ -119,12 +118,20 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { // Register the functions as a submodule let funcs = PyModule::new(py, "functions")?; functions::init_module(&funcs)?; + // Spark-compatible functions live under `functions.spark`. + let spark_funcs = PyModule::new(py, "spark")?; + spark_functions::init_module(&spark_funcs)?; + funcs.add_submodule(&spark_funcs)?; m.add_submodule(&funcs)?; let store = PyModule::new(py, "object_store")?; store::init_module(&store)?; m.add_submodule(&store)?; + let options = PyModule::new(py, "options")?; + options::init_module(&options)?; + m.add_submodule(&options)?; + // Register substrait as a submodule #[cfg(feature = "substrait")] setup_substrait_module(py, &m)?; 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/options.rs b/crates/core/src/options.rs new file mode 100644 index 000000000..6b6037695 --- /dev/null +++ b/crates/core/src/options.rs @@ -0,0 +1,159 @@ +// 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 arrow::datatypes::{DataType, Schema}; +use arrow::pyarrow::PyArrowType; +use datafusion::prelude::CsvReadOptions; +use pyo3::prelude::{PyModule, PyModuleMethods}; +use pyo3::{Bound, PyResult, pyclass, pymethods}; + +use crate::context::parse_file_compression_type; +use crate::errors::PyDataFusionError; +use crate::expr::sort_expr::PySortExpr; + +/// Options for reading CSV files +#[pyclass(name = "CsvReadOptions", module = "datafusion.options", frozen)] +pub struct PyCsvReadOptions { + pub has_header: bool, + pub delimiter: u8, + pub quote: u8, + pub terminator: Option, + pub escape: Option, + pub comment: Option, + pub newlines_in_values: bool, + pub schema: Option>, + pub schema_infer_max_records: usize, + pub file_extension: String, + pub table_partition_cols: Vec<(String, PyArrowType)>, + pub file_compression_type: String, + pub file_sort_order: Vec>, + pub null_regex: Option, + pub truncated_rows: bool, +} + +#[pymethods] +impl PyCsvReadOptions { + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = ( + has_header=true, + delimiter=b',', + quote=b'"', + terminator=None, + escape=None, + comment=None, + newlines_in_values=false, + schema=None, + schema_infer_max_records=1000, + file_extension=".csv".to_string(), + table_partition_cols=vec![], + file_compression_type="".to_string(), + file_sort_order=vec![], + null_regex=None, + truncated_rows=false + ))] + #[new] + fn new( + has_header: bool, + delimiter: u8, + quote: u8, + terminator: Option, + escape: Option, + comment: Option, + newlines_in_values: bool, + schema: Option>, + schema_infer_max_records: usize, + file_extension: String, + table_partition_cols: Vec<(String, PyArrowType)>, + file_compression_type: String, + file_sort_order: Vec>, + null_regex: Option, + truncated_rows: bool, + ) -> Self { + Self { + has_header, + delimiter, + quote, + terminator, + escape, + comment, + newlines_in_values, + schema, + schema_infer_max_records, + file_extension, + table_partition_cols, + file_compression_type, + file_sort_order, + null_regex, + truncated_rows, + } + } +} + +impl<'a> TryFrom<&'a PyCsvReadOptions> for CsvReadOptions<'a> { + type Error = PyDataFusionError; + + fn try_from(value: &'a PyCsvReadOptions) -> Result, Self::Error> { + let partition_cols: Vec<(String, DataType)> = value + .table_partition_cols + .iter() + .map(|(name, dtype)| (name.clone(), dtype.0.clone())) + .collect(); + + let compression = parse_file_compression_type(Some(value.file_compression_type.clone()))?; + + let sort_order: Vec> = value + .file_sort_order + .iter() + .map(|inner| { + inner + .iter() + .map(|sort_expr| sort_expr.sort.clone()) + .collect() + }) + .collect(); + + // Explicit struct initialization to catch upstream changes + let mut options = CsvReadOptions { + has_header: value.has_header, + delimiter: value.delimiter, + quote: value.quote, + terminator: value.terminator, + escape: value.escape, + comment: value.comment, + newlines_in_values: value.newlines_in_values, + schema: None, // Will be set separately due to lifetime constraints + schema_infer_max_records: value.schema_infer_max_records, + file_extension: value.file_extension.as_str(), + table_partition_cols: partition_cols, + file_compression_type: compression, + file_sort_order: sort_order, + null_regex: value.null_regex.clone(), + truncated_rows: value.truncated_rows, + }; + + // Set schema separately to handle the lifetime + options.schema = value.schema.as_ref().map(|s| &s.0); + + Ok(options) + } +} + +pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + + Ok(()) +} diff --git a/src/physical_plan.rs b/crates/core/src/physical_plan.rs similarity index 65% rename from src/physical_plan.rs rename to crates/core/src/physical_plan.rs index 645649e2c..594655a60 100644 --- a/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -17,17 +17,25 @@ use std::sync::Arc; -use datafusion::physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties}; -use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; +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; - -#[pyclass(frozen, name = "ExecutionPlan", module = "datafusion", subclass)] +use crate::metrics::PyMetricsSet; + +#[pyclass( + from_py_object, + frozen, + name = "ExecutionPlan", + module = "datafusion", + subclass +)] #[derive(Debug, Clone)] pub struct PyExecutionPlan { pub plan: Arc, @@ -61,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(); @@ -73,23 +96,28 @@ impl PyExecutionPlan { } #[staticmethod] - pub fn from_proto( + pub fn from_bytes( ctx: PySessionContext, proto_msg: Bound<'_, PyBytes>, ) -> PyDataFusionResult { - let bytes: &[u8] = proto_msg.extract()?; + let bytes: &[u8] = proto_msg.extract().map_err(Into::::into)?; 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)) } + pub fn metrics(&self) -> Option { + self.plan.metrics().map(PyMetricsSet::new) + } + fn __repr__(&self) -> String { self.display_indent() } diff --git a/src/pyarrow_filter_expression.rs b/crates/core/src/pyarrow_filter_expression.rs similarity index 98% rename from src/pyarrow_filter_expression.rs rename to crates/core/src/pyarrow_filter_expression.rs index c9d3df32d..e3b4b6009 100644 --- a/src/pyarrow_filter_expression.rs +++ b/crates/core/src/pyarrow_filter_expression.rs @@ -22,7 +22,7 @@ use datafusion::common::{Column, ScalarValue}; use datafusion::logical_expr::expr::InList; use datafusion::logical_expr::{Between, BinaryExpr, Expr, Operator}; /// Converts a Datafusion logical plan expression (Expr) into a PyArrow compute expression -use pyo3::{prelude::*, IntoPyObjectExt}; +use pyo3::{IntoPyObjectExt, prelude::*}; use crate::errors::{PyDataFusionError, PyDataFusionResult}; use crate::pyarrow_util::scalar_to_pyarrow; @@ -47,7 +47,7 @@ fn operator_to_py<'py>( _ => { return Err(PyDataFusionError::Common(format!( "Unsupported operator {operator:?}" - ))) + ))); } }; Ok(py_op) @@ -57,7 +57,7 @@ fn extract_scalar_list<'py>( exprs: &[Expr], py: Python<'py>, ) -> PyDataFusionResult>> { - let ret = exprs + exprs .iter() .map(|expr| match expr { // TODO: should we also leverage `ScalarValue::to_pyarrow` here? @@ -83,8 +83,7 @@ fn extract_scalar_list<'py>( "Only a list of Literals are supported got {expr:?}" ))), }) - .collect(); - ret + .collect() } impl PyArrowFilterExpression { diff --git a/crates/core/src/pyarrow_util.rs b/crates/core/src/pyarrow_util.rs new file mode 100644 index 000000000..1401a4938 --- /dev/null +++ b/crates/core/src/pyarrow_util.rs @@ -0,0 +1,163 @@ +// 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. + +//! Conversions between PyArrow and DataFusion types + +use std::sync::Arc; + +use arrow::array::{Array, ArrayData, ArrayRef, ListArray, make_array}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::Field; +use arrow::pyarrow::{FromPyArrow, ToPyArrow}; +use datafusion::common::exec_err; +use datafusion::scalar::ScalarValue; +use pyo3::types::{PyAnyMethods, PyList}; +use pyo3::{Borrowed, Bound, FromPyObject, PyAny, PyErr, PyResult, Python}; + +use crate::common::data_type::PyScalarValue; +use crate::errors::PyDataFusionError; + +/// Helper function to turn an Array into a ScalarValue. If ``as_list_array`` is true, +/// the array will be turned into a ``ListArray``. Otherwise, we extract the first value +/// from the array. +fn array_to_scalar_value(array: ArrayRef, as_list_array: bool) -> PyResult { + if as_list_array { + let field = Arc::new(Field::new_list_field( + array.data_type().clone(), + array.nulls().is_some(), + )); + let offsets = OffsetBuffer::from_lengths(vec![array.len()]); + let list_array = ListArray::new(field, offsets, array, None); + Ok(PyScalarValue(ScalarValue::List(Arc::new(list_array)))) + } else { + let scalar = ScalarValue::try_from_array(&array, 0).map_err(PyDataFusionError::from)?; + Ok(PyScalarValue(scalar)) + } +} + +/// Helper function to take any Python object that contains an Arrow PyCapsule +/// interface and attempt to extract a scalar value from it. If `as_list_array` +/// is true, the array will be turned into a ``ListArray``. Otherwise, we extract +/// the first value from the array. +fn pyobj_extract_scalar_via_capsule( + value: &Bound<'_, PyAny>, + as_list_array: bool, +) -> PyResult { + let array_data = ArrayData::from_pyarrow_bound(value)?; + let array = make_array(array_data); + + array_to_scalar_value(array, as_list_array) +} + +impl FromPyArrow for PyScalarValue { + fn from_pyarrow_bound(value: &Bound<'_, PyAny>) -> PyResult { + let py = value.py(); + let pyarrow_mod = py.import("pyarrow"); + + // Is it a PyArrow object? + if let Ok(pa) = pyarrow_mod.as_ref() { + let scalar_type = pa.getattr("Scalar")?; + if value.is_instance(&scalar_type)? { + let typ = value.getattr("type")?; + + // construct pyarrow array from the python value and pyarrow type + let factory = py.import("pyarrow")?.getattr("array")?; + let args = PyList::new(py, [value])?; + let array = factory.call1((args, typ))?; + + return pyobj_extract_scalar_via_capsule(&array, false); + } + + let array_type = pa.getattr("Array")?; + if value.is_instance(&array_type)? { + return pyobj_extract_scalar_via_capsule(value, true); + } + } + + // Is it a NanoArrow scalar? + if let Ok(na) = py.import("nanoarrow") { + let scalar_type = py.import("nanoarrow.array")?.getattr("Scalar")?; + if value.is_instance(&scalar_type)? { + return pyobj_extract_scalar_via_capsule(value, false); + } + let array_type = na.getattr("Array")?; + if value.is_instance(&array_type)? { + return pyobj_extract_scalar_via_capsule(value, true); + } + } + + // Is it a arro3 scalar? + if let Ok(arro3) = py.import("arro3").and_then(|arro3| arro3.getattr("core")) { + let scalar_type = arro3.getattr("Scalar")?; + if value.is_instance(&scalar_type)? { + return pyobj_extract_scalar_via_capsule(value, false); + } + let array_type = arro3.getattr("Array")?; + if value.is_instance(&array_type)? { + return pyobj_extract_scalar_via_capsule(value, true); + } + } + + // Does it have a PyCapsule interface but isn't one of our known libraries? + // If so do our "best guess". Try checking type name, and if that fails + // return a single value if the length is 1 and return a List value otherwise + if value.hasattr("__arrow_c_array__")? { + let type_name = value.get_type().repr()?; + if type_name.contains("Scalar")? { + return pyobj_extract_scalar_via_capsule(value, false); + } + if type_name.contains("Array")? { + return pyobj_extract_scalar_via_capsule(value, true); + } + + let array_data = ArrayData::from_pyarrow_bound(value)?; + let array = make_array(array_data); + + let as_array_list = array.len() != 1; + return array_to_scalar_value(array, as_array_list); + } + + // Last attempt - try to create a PyArrow scalar from a plain Python object + if let Ok(pa) = pyarrow_mod.as_ref() { + let scalar = pa.call_method1("scalar", (value,))?; + + PyScalarValue::from_pyarrow_bound(&scalar) + } else { + exec_err!("Unable to import scalar value").map_err(PyDataFusionError::from)? + } + } +} + +impl<'source> FromPyObject<'_, 'source> for PyScalarValue { + type Error = PyErr; + + fn extract(value: Borrowed<'_, 'source, PyAny>) -> Result { + Self::from_pyarrow_bound(&value) + } +} + +pub fn scalar_to_pyarrow<'py>( + scalar: &ScalarValue, + py: Python<'py>, +) -> PyResult> { + let array = scalar.to_array().map_err(PyDataFusionError::from)?; + // convert to pyarrow array using C data interface + let pyarray = array.to_data().to_pyarrow(py)?; + let pyscalar = pyarray.call_method1("__getitem__", (0,))?; + + Ok(pyscalar) +} diff --git a/src/record_batch.rs b/crates/core/src/record_batch.rs similarity index 97% rename from src/record_batch.rs rename to crates/core/src/record_batch.rs index 2e50ba75e..0492c6c76 100644 --- a/src/record_batch.rs +++ b/crates/core/src/record_batch.rs @@ -20,14 +20,14 @@ use std::sync::Arc; use datafusion::arrow::pyarrow::ToPyArrow; use datafusion::arrow::record_batch::RecordBatch; use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion_python_util::wait_for_future; use futures::StreamExt; use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration}; use pyo3::prelude::*; -use pyo3::{pyclass, pymethods, PyAny, PyResult, Python}; +use pyo3::{PyAny, PyResult, Python, pyclass, pymethods}; use tokio::sync::Mutex; use crate::errors::PyDataFusionError; -use crate::utils::wait_for_future; #[pyclass(name = "RecordBatch", module = "datafusion", subclass, frozen)] pub struct PyRecordBatch { diff --git a/crates/core/src/spark_functions.rs b/crates/core/src/spark_functions.rs new file mode 100644 index 000000000..e7cb94f8c --- /dev/null +++ b/crates/core/src/spark_functions.rs @@ -0,0 +1,363 @@ +// 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. + +//! PyO3 wrappers for the [`datafusion-spark`] crate. +//! +//! Exposes Spark-compatible scalar and aggregate function builders for use +//! from Python under `datafusion.functions.spark`. + +use datafusion::logical_expr::Expr; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion_spark::{expr_fn, function as udf}; +use pyo3::prelude::*; +use pyo3::wrap_pyfunction; + +use crate::common::data_type::NullTreatment; +use crate::errors::PyDataFusionResult; +use crate::expr::PyExpr; +use crate::expr::sort_expr::PySortExpr; +use crate::functions::add_builder_fns_to_aggregate; + +/// Generates a [pyo3] wrapper for [datafusion_spark::expr_fn]. +/// +/// These functions have explicit named arguments and mirror the upstream +/// `expr_fn::$FUNC` signature. +macro_rules! spark_expr_fn { + ($FUNC:ident) => { + spark_expr_fn!($FUNC,); + }; + ($FUNC:ident, $($arg:ident)*) => { + #[pyfunction] + fn $FUNC($($arg: PyExpr),*) -> PyExpr { + expr_fn::$FUNC($($arg.into()),*).into() + } + }; +} + +/// Generates a variadic [pyo3] wrapper that calls the [`ScalarUDF`] factory +/// directly. Required for functions whose upstream `expr_fn` wrapper accepts +/// a single `Expr` instead of `Vec` (an upstream `export_functions!` +/// macro-arm quirk), so we bypass it to get true Python `*args` semantics. +macro_rules! spark_udf_vec { + ($PY_NAME:ident, $UDF_PATH:path) => { + #[pyfunction] + #[pyo3(signature = (*args))] + fn $PY_NAME(args: Vec) -> PyExpr { + let udf = $UDF_PATH(); + let args: Vec = args.into_iter().map(Into::into).collect(); + Expr::ScalarFunction(ScalarFunction::new_udf(udf, args)).into() + } + }; +} + +/// Generates a [pyo3] wrapper for Spark aggregate functions. Mirrors +/// [`crate::functions::aggregate_function`] but points at +/// [`datafusion_spark::expr_fn`]. +macro_rules! spark_aggregate { + ($NAME:ident) => { + spark_aggregate!($NAME, expr); + }; + ($NAME:ident, $($arg:ident)*) => { + #[pyfunction] + #[pyo3(signature = ($($arg),*, distinct=None, filter=None, order_by=None, null_treatment=None))] + fn $NAME( + $($arg: PyExpr),*, + distinct: Option, + filter: Option, + order_by: Option>, + null_treatment: Option, + ) -> PyDataFusionResult { + let agg_fn = expr_fn::$NAME($($arg.into()),*); + add_builder_fns_to_aggregate(agg_fn, distinct, filter, order_by, null_treatment) + } + }; +} + +// --------------------------------------------------------------------------- +// Aggregate functions +// --------------------------------------------------------------------------- + +spark_aggregate!(avg, arg1); +spark_aggregate!(try_sum, arg1); +spark_aggregate!(collect_list, arg1); +spark_aggregate!(collect_set, arg1); + +// --------------------------------------------------------------------------- +// Array functions +// --------------------------------------------------------------------------- + +// Upstream factory is `spark_array_contains`; expose under the Spark SQL +// name `array_contains` on the Python side. +#[pyfunction] +fn array_contains(arr: PyExpr, element: PyExpr) -> PyExpr { + expr_fn::spark_array_contains(arr.into(), element.into()).into() +} +spark_udf_vec!(array, udf::array::array); +spark_expr_fn!(shuffle, arg1); +spark_expr_fn!(array_repeat, element count); +spark_expr_fn!(slice, arr start length); + +// --------------------------------------------------------------------------- +// Bitmap functions +// --------------------------------------------------------------------------- + +spark_expr_fn!(bitmap_count, arg1); +spark_expr_fn!(bitmap_bit_position, arg1); +spark_expr_fn!(bitmap_bucket_number, arg1); + +// --------------------------------------------------------------------------- +// Bitwise functions +// --------------------------------------------------------------------------- + +spark_expr_fn!(bit_get, col pos); +spark_expr_fn!(bit_count, col); +spark_expr_fn!(bitwise_not, col); +spark_expr_fn!(shiftleft, value shift); +spark_expr_fn!(shiftright, value shift); +spark_expr_fn!(shiftrightunsigned, value shift); + +// --------------------------------------------------------------------------- +// Collection / Conditional / Conversion +// --------------------------------------------------------------------------- + +spark_expr_fn!(size, arg1); + +// Python keyword `if` → exposed as `if_`. Upstream Rust ident is `r#if`. +#[pyfunction] +fn if_(condition: PyExpr, if_true: PyExpr, if_false: PyExpr) -> PyExpr { + expr_fn::r#if(condition.into(), if_true.into(), if_false.into()).into() +} + +// `spark_cast` is config-injected by the upstream `expr_fn` helper; defaults +// applied automatically there. +spark_expr_fn!(spark_cast, arg1 arg2); + +// --------------------------------------------------------------------------- +// Datetime functions +// --------------------------------------------------------------------------- + +spark_expr_fn!(add_months, start_date num_months); +spark_expr_fn!(date_add, start_date days); +spark_expr_fn!(date_sub, start_date days); +spark_expr_fn!(hour, arg1); +spark_expr_fn!(minute, arg1); +spark_expr_fn!(second, arg1); +spark_expr_fn!(last_day, arg1); +spark_expr_fn!(make_dt_interval, days hours mins secs); +spark_expr_fn!(make_interval, years months weeks days hours mins secs); +spark_expr_fn!(next_day, start_date day_of_week); +spark_expr_fn!(date_diff, end_date start_date); +spark_expr_fn!(date_trunc, fmt ts); +spark_expr_fn!(time_trunc, fmt t); +spark_expr_fn!(trunc, dt fmt); +spark_expr_fn!(date_part, field source); +spark_expr_fn!(from_utc_timestamp, ts tz); +spark_expr_fn!(to_utc_timestamp, ts tz); +spark_expr_fn!(unix_date, dt); +spark_expr_fn!(unix_micros, ts); +spark_expr_fn!(unix_millis, ts); +spark_expr_fn!(unix_seconds, ts); + +// --------------------------------------------------------------------------- +// Hash functions +// --------------------------------------------------------------------------- + +spark_expr_fn!(crc32, arg1); +spark_expr_fn!(sha1, arg1); +spark_expr_fn!(sha2, arg1 bit_length); +spark_udf_vec!(xxhash64, udf::hash::xxhash64); + +// --------------------------------------------------------------------------- +// JSON functions +// --------------------------------------------------------------------------- + +spark_udf_vec!(json_tuple, udf::json::json_tuple); + +// --------------------------------------------------------------------------- +// Map functions +// --------------------------------------------------------------------------- + +spark_expr_fn!(map_from_arrays, keys values); +spark_expr_fn!(map_from_entries, arg1); +spark_expr_fn!(str_to_map, text pair_delim key_value_delim); + +// --------------------------------------------------------------------------- +// Math functions +// --------------------------------------------------------------------------- + +spark_expr_fn!(abs, arg1); +spark_expr_fn!(ceil, arg1); +spark_expr_fn!(expm1, arg1); +spark_expr_fn!(factorial, arg1); +spark_expr_fn!(floor, arg1); +spark_expr_fn!(hex, arg1); +spark_expr_fn!(modulus, dividend divisor); +spark_expr_fn!(pmod, dividend divisor); +spark_expr_fn!(rint, arg1); +spark_expr_fn!(round, value scale); +spark_expr_fn!(unhex, arg1); +spark_expr_fn!(width_bucket, value min_value max_value num_buckets); +spark_expr_fn!(csc, arg1); +spark_expr_fn!(sec, arg1); +spark_expr_fn!(negative, arg1); +spark_expr_fn!(bin, arg1); + +// --------------------------------------------------------------------------- +// String functions +// --------------------------------------------------------------------------- + +spark_expr_fn!(ascii, arg1); +spark_expr_fn!(base64, bin_input); +// `char` collides with the Rust primitive type in macro hygiene; rename the +// Rust ident and re-expose under the original name to Python. +#[pyfunction] +#[pyo3(name = "char")] +fn char_fn(arg1: PyExpr) -> PyExpr { + expr_fn::char(arg1.into()).into() +} +spark_udf_vec!(concat, udf::string::concat); +spark_udf_vec!(elt, udf::string::elt); +spark_expr_fn!(ilike, str pattern); +spark_expr_fn!(length, arg1); +spark_expr_fn!(like, str pattern); +spark_expr_fn!(luhn_check, arg1); +spark_udf_vec!(format_string, udf::string::format_string); +spark_expr_fn!(space, arg1); +spark_expr_fn!(substring, str pos length); +spark_expr_fn!(unbase64, str); +spark_expr_fn!(soundex, str); +spark_expr_fn!(is_valid_utf8, str); +spark_expr_fn!(make_valid_utf8, str); + +// --------------------------------------------------------------------------- +// URL functions +// --------------------------------------------------------------------------- + +spark_udf_vec!(parse_url, udf::url::parse_url); +spark_udf_vec!(try_parse_url, udf::url::try_parse_url); +spark_udf_vec!(url_decode, udf::url::url_decode); +spark_udf_vec!(try_url_decode, udf::url::try_url_decode); +spark_udf_vec!(url_encode, udf::url::url_encode); + +// --------------------------------------------------------------------------- +// Module init +// --------------------------------------------------------------------------- + +pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + // Aggregate + m.add_wrapped(wrap_pyfunction!(avg))?; + m.add_wrapped(wrap_pyfunction!(try_sum))?; + m.add_wrapped(wrap_pyfunction!(collect_list))?; + m.add_wrapped(wrap_pyfunction!(collect_set))?; + // Array + m.add_wrapped(wrap_pyfunction!(array_contains))?; + m.add_wrapped(wrap_pyfunction!(array))?; + m.add_wrapped(wrap_pyfunction!(shuffle))?; + m.add_wrapped(wrap_pyfunction!(array_repeat))?; + m.add_wrapped(wrap_pyfunction!(slice))?; + // Bitmap + m.add_wrapped(wrap_pyfunction!(bitmap_count))?; + m.add_wrapped(wrap_pyfunction!(bitmap_bit_position))?; + m.add_wrapped(wrap_pyfunction!(bitmap_bucket_number))?; + // Bitwise + m.add_wrapped(wrap_pyfunction!(bit_get))?; + m.add_wrapped(wrap_pyfunction!(bit_count))?; + m.add_wrapped(wrap_pyfunction!(bitwise_not))?; + m.add_wrapped(wrap_pyfunction!(shiftleft))?; + m.add_wrapped(wrap_pyfunction!(shiftright))?; + m.add_wrapped(wrap_pyfunction!(shiftrightunsigned))?; + // Collection + m.add_wrapped(wrap_pyfunction!(size))?; + // Conditional + m.add_wrapped(wrap_pyfunction!(if_))?; + // Conversion + m.add_wrapped(wrap_pyfunction!(spark_cast))?; + // Datetime + m.add_wrapped(wrap_pyfunction!(add_months))?; + m.add_wrapped(wrap_pyfunction!(date_add))?; + m.add_wrapped(wrap_pyfunction!(date_sub))?; + m.add_wrapped(wrap_pyfunction!(hour))?; + m.add_wrapped(wrap_pyfunction!(minute))?; + m.add_wrapped(wrap_pyfunction!(second))?; + m.add_wrapped(wrap_pyfunction!(last_day))?; + m.add_wrapped(wrap_pyfunction!(make_dt_interval))?; + m.add_wrapped(wrap_pyfunction!(make_interval))?; + m.add_wrapped(wrap_pyfunction!(next_day))?; + m.add_wrapped(wrap_pyfunction!(date_diff))?; + m.add_wrapped(wrap_pyfunction!(date_trunc))?; + m.add_wrapped(wrap_pyfunction!(time_trunc))?; + m.add_wrapped(wrap_pyfunction!(trunc))?; + m.add_wrapped(wrap_pyfunction!(date_part))?; + m.add_wrapped(wrap_pyfunction!(from_utc_timestamp))?; + m.add_wrapped(wrap_pyfunction!(to_utc_timestamp))?; + m.add_wrapped(wrap_pyfunction!(unix_date))?; + m.add_wrapped(wrap_pyfunction!(unix_micros))?; + m.add_wrapped(wrap_pyfunction!(unix_millis))?; + m.add_wrapped(wrap_pyfunction!(unix_seconds))?; + // Hash + m.add_wrapped(wrap_pyfunction!(crc32))?; + m.add_wrapped(wrap_pyfunction!(sha1))?; + m.add_wrapped(wrap_pyfunction!(sha2))?; + m.add_wrapped(wrap_pyfunction!(xxhash64))?; + // JSON + m.add_wrapped(wrap_pyfunction!(json_tuple))?; + // Map + m.add_wrapped(wrap_pyfunction!(map_from_arrays))?; + m.add_wrapped(wrap_pyfunction!(map_from_entries))?; + m.add_wrapped(wrap_pyfunction!(str_to_map))?; + // Math + m.add_wrapped(wrap_pyfunction!(abs))?; + m.add_wrapped(wrap_pyfunction!(ceil))?; + m.add_wrapped(wrap_pyfunction!(expm1))?; + m.add_wrapped(wrap_pyfunction!(factorial))?; + m.add_wrapped(wrap_pyfunction!(floor))?; + m.add_wrapped(wrap_pyfunction!(hex))?; + m.add_wrapped(wrap_pyfunction!(modulus))?; + m.add_wrapped(wrap_pyfunction!(pmod))?; + m.add_wrapped(wrap_pyfunction!(rint))?; + m.add_wrapped(wrap_pyfunction!(round))?; + m.add_wrapped(wrap_pyfunction!(unhex))?; + m.add_wrapped(wrap_pyfunction!(width_bucket))?; + m.add_wrapped(wrap_pyfunction!(csc))?; + m.add_wrapped(wrap_pyfunction!(sec))?; + m.add_wrapped(wrap_pyfunction!(negative))?; + m.add_wrapped(wrap_pyfunction!(bin))?; + // String + m.add_wrapped(wrap_pyfunction!(ascii))?; + m.add_wrapped(wrap_pyfunction!(base64))?; + m.add_wrapped(wrap_pyfunction!(char_fn))?; + m.add_wrapped(wrap_pyfunction!(concat))?; + m.add_wrapped(wrap_pyfunction!(elt))?; + m.add_wrapped(wrap_pyfunction!(ilike))?; + m.add_wrapped(wrap_pyfunction!(length))?; + m.add_wrapped(wrap_pyfunction!(like))?; + m.add_wrapped(wrap_pyfunction!(luhn_check))?; + m.add_wrapped(wrap_pyfunction!(format_string))?; + m.add_wrapped(wrap_pyfunction!(space))?; + m.add_wrapped(wrap_pyfunction!(substring))?; + m.add_wrapped(wrap_pyfunction!(unbase64))?; + m.add_wrapped(wrap_pyfunction!(soundex))?; + m.add_wrapped(wrap_pyfunction!(is_valid_utf8))?; + m.add_wrapped(wrap_pyfunction!(make_valid_utf8))?; + // URL + m.add_wrapped(wrap_pyfunction!(parse_url))?; + m.add_wrapped(wrap_pyfunction!(try_parse_url))?; + m.add_wrapped(wrap_pyfunction!(url_decode))?; + m.add_wrapped(wrap_pyfunction!(try_url_decode))?; + m.add_wrapped(wrap_pyfunction!(url_encode))?; + Ok(()) +} diff --git a/src/sql.rs b/crates/core/src/sql.rs similarity index 100% rename from src/sql.rs rename to crates/core/src/sql.rs diff --git a/src/sql/exceptions.rs b/crates/core/src/sql/exceptions.rs similarity index 100% rename from src/sql/exceptions.rs rename to crates/core/src/sql/exceptions.rs diff --git a/src/sql/logical.rs b/crates/core/src/sql/logical.rs similarity index 86% rename from src/sql/logical.rs rename to crates/core/src/sql/logical.rs index 37f20d287..12fc43bdc 100644 --- a/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; @@ -55,7 +56,8 @@ use crate::expr::recursive_query::PyRecursiveQuery; use crate::expr::repartition::PyRepartition; use crate::expr::sort::PySort; use crate::expr::statement::{ - PyDeallocate, PyExecute, PyPrepare, PySetVariable, PyTransactionEnd, PyTransactionStart, + PyDeallocate, PyExecute, PyPrepare, PyResetVariable, PySetVariable, PyTransactionEnd, + PyTransactionStart, }; use crate::expr::subquery::PySubquery; use crate::expr::subquery_alias::PySubqueryAlias; @@ -65,8 +67,15 @@ use crate::expr::unnest::PyUnnest; use crate::expr::values::PyValues; use crate::expr::window::PyWindowExpr; -#[pyclass(frozen, name = "LogicalPlan", module = "datafusion", subclass)] -#[derive(Debug, Clone)] +#[pyclass( + from_py_object, + frozen, + name = "LogicalPlan", + module = "datafusion", + subclass, + eq +)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct PyLogicalPlan { pub(crate) plan: Arc, } @@ -115,6 +124,9 @@ impl PyLogicalPlan { PyTransactionEnd::from(plan.clone()).to_variant(py) } Statement::SetVariable(plan) => PySetVariable::from(plan.clone()).to_variant(py), + Statement::ResetVariable(plan) => { + PyResetVariable::from(plan.clone()).to_variant(py) + } Statement::Prepare(plan) => PyPrepare::from(plan.clone()).to_variant(py), Statement::Execute(plan) => PyExecute::from(plan.clone()).to_variant(py), Statement::Deallocate(plan) => PyDeallocate::from(plan.clone()).to_variant(py), @@ -123,7 +135,7 @@ impl PyLogicalPlan { LogicalPlan::Dml(plan) => PyDmlStatement::from(plan.clone()).to_variant(py), LogicalPlan::Ddl(plan) => match plan { DdlStatement::CreateExternalTable(plan) => { - PyCreateExternalTable::from(plan.clone()).to_variant(py) + PyCreateExternalTable::from(plan.as_ref().clone()).to_variant(py) } DdlStatement::CreateMemoryTable(plan) => { PyCreateMemoryTable::from(plan.clone()).to_variant(py) @@ -142,7 +154,7 @@ impl PyLogicalPlan { PyDropCatalogSchema::from(plan.clone()).to_variant(py) } DdlStatement::CreateFunction(plan) => { - PyCreateFunction::from(plan.clone()).to_variant(py) + PyCreateFunction::from(plan.as_ref().clone()).to_variant(py) } DdlStatement::DropFunction(plan) => { PyDropFunction::from(plan.clone()).to_variant(py) @@ -185,21 +197,33 @@ 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 { - let bytes: &[u8] = proto_msg.extract()?; + let bytes: &[u8] = proto_msg.extract().map_err(Into::::into)?; let proto_plan = datafusion_proto::protobuf::LogicalPlanNode::decode(bytes).map_err(|e| { PyRuntimeError::new_err(format!( @@ -207,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/src/sql/util.rs b/crates/core/src/sql/util.rs similarity index 97% rename from src/sql/util.rs rename to crates/core/src/sql/util.rs index 5edff006f..d1e8964f8 100644 --- a/src/sql/util.rs +++ b/crates/core/src/sql/util.rs @@ -17,7 +17,7 @@ use std::collections::HashMap; -use datafusion::common::{exec_err, plan_datafusion_err, DataFusionError}; +use datafusion::common::{DataFusionError, exec_err, plan_datafusion_err}; use datafusion::logical_expr::sqlparser::dialect::dialect_from_str; use datafusion::sql::sqlparser::dialect::Dialect; use datafusion::sql::sqlparser::parser::Parser; diff --git a/src/store.rs b/crates/core/src/store.rs similarity index 94% rename from src/store.rs rename to crates/core/src/store.rs index 3eae866bc..8535e83b7 100644 --- a/src/store.rs +++ b/crates/core/src/store.rs @@ -36,6 +36,7 @@ pub enum StorageContexts { } #[pyclass( + from_py_object, frozen, name = "LocalFileSystem", module = "datafusion.store", @@ -66,7 +67,13 @@ impl PyLocalFileSystemContext { } } -#[pyclass(frozen, name = "MicrosoftAzure", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "MicrosoftAzure", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyMicrosoftAzureContext { pub inner: Arc, @@ -143,7 +150,13 @@ impl PyMicrosoftAzureContext { } } -#[pyclass(frozen, name = "GoogleCloud", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "GoogleCloud", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyGoogleCloudContext { pub inner: Arc, @@ -173,7 +186,13 @@ impl PyGoogleCloudContext { } } -#[pyclass(frozen, name = "AmazonS3", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "AmazonS3", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyAmazonS3Context { pub inner: Arc, @@ -237,7 +256,13 @@ impl PyAmazonS3Context { } } -#[pyclass(frozen, name = "Http", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Http", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyHttpContext { pub url: String, diff --git a/src/substrait.rs b/crates/core/src/substrait.rs similarity index 77% rename from src/substrait.rs rename to crates/core/src/substrait.rs index 7b06aff74..dcb5587cf 100644 --- a/src/substrait.rs +++ b/crates/core/src/substrait.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use datafusion_python_util::wait_for_future; use datafusion_substrait::logical_plan::{consumer, producer}; use datafusion_substrait::serializer; use datafusion_substrait::substrait::proto::Plan; @@ -23,11 +24,16 @@ use pyo3::prelude::*; use pyo3::types::PyBytes; use crate::context::PySessionContext; -use crate::errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult}; +use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err}; use crate::sql::logical::PyLogicalPlan; -use crate::utils::wait_for_future; -#[pyclass(frozen, name = "Plan", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Plan", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PyPlan { pub plan: Plan, @@ -42,6 +48,19 @@ impl PyPlan { .map_err(PyDataFusionError::EncodeError)?; Ok(PyBytes::new(py, &proto_bytes).into()) } + + /// Get the JSON representation of the substrait plan + fn to_json(&self) -> PyDataFusionResult { + let json = serde_json::to_string_pretty(&self.plan).map_err(to_datafusion_err)?; + Ok(json) + } + + /// Parse a Substrait Plan from its JSON representation + #[staticmethod] + fn from_json(json: &str) -> PyDataFusionResult { + let plan: Plan = serde_json::from_str(json).map_err(to_datafusion_err)?; + Ok(PyPlan { plan }) + } } impl From for Plan { @@ -59,7 +78,13 @@ impl From for PyPlan { /// A PySubstraitSerializer is a representation of a Serializer that is capable of both serializing /// a `LogicalPlan` instance to Substrait Protobuf bytes and also deserialize Substrait Protobuf bytes /// to a valid `LogicalPlan` instance. -#[pyclass(frozen, name = "Serde", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Serde", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PySubstraitSerializer; @@ -83,8 +108,8 @@ impl PySubstraitSerializer { py: Python, ) -> PyDataFusionResult { PySubstraitSerializer::serialize_bytes(sql, ctx, py).and_then(|proto_bytes| { - let proto_bytes = proto_bytes.bind(py).downcast::().unwrap(); - PySubstraitSerializer::deserialize_bytes(proto_bytes.as_bytes().to_vec(), py) + let proto_bytes = proto_bytes.bind(py).cast::().unwrap(); + PySubstraitSerializer::deserialize_bytes(proto_bytes.as_bytes().to_vec()) }) } @@ -106,13 +131,19 @@ impl PySubstraitSerializer { } #[staticmethod] - pub fn deserialize_bytes(proto_bytes: Vec, py: Python) -> PyDataFusionResult { - let plan = wait_for_future(py, serializer::deserialize_bytes(proto_bytes))??; + pub fn deserialize_bytes(proto_bytes: Vec) -> PyDataFusionResult { + let plan = serializer::deserialize_bytes(&proto_bytes)?; Ok(PyPlan { plan: *plan }) } } -#[pyclass(frozen, name = "Producer", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Producer", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PySubstraitProducer; @@ -129,7 +160,13 @@ impl PySubstraitProducer { } } -#[pyclass(frozen, name = "Consumer", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Consumer", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PySubstraitConsumer; diff --git a/src/table.rs b/crates/core/src/table.rs similarity index 67% rename from src/table.rs rename to crates/core/src/table.rs index 0eec57f75..e0f0f0d13 100644 --- a/src/table.rs +++ b/crates/core/src/table.rs @@ -15,28 +15,40 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::pyarrow::ToPyArrow; use async_trait::async_trait; -use datafusion::catalog::Session; +use datafusion::catalog::{Session, TableProviderFactory}; use datafusion::common::Column; use datafusion::datasource::{TableProvider, TableType}; -use datafusion::logical_expr::{Expr, LogicalPlanBuilder, TableProviderFilterPushDown}; +use datafusion::logical_expr::{ + CreateExternalTable, Expr, LogicalPlanBuilder, TableProviderFilterPushDown, +}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::DataFrame; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_python_util::{create_logical_extension_capsule, table_provider_from_pycapsule}; +use pyo3::IntoPyObjectExt; use pyo3::prelude::*; +use crate::context::PySessionContext; use crate::dataframe::PyDataFrame; use crate::dataset::Dataset; -use crate::utils::table_provider_from_pycapsule; +use crate::errors; +use crate::expr::create_external_table::PyCreateExternalTable; /// This struct is used as a common method for all TableProviders, /// whether they refer to an FFI provider, an internally known /// implementation, a dataset, or a dataframe view. -#[pyclass(frozen, name = "RawTable", module = "datafusion.catalog", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "RawTable", + module = "datafusion.catalog", + subclass +)] #[derive(Clone)] pub struct PyTable { pub table: Arc, @@ -60,12 +72,13 @@ impl PyTable { /// - FFI Table Providers via PyCapsule /// - PyArrow Dataset objects #[new] - pub fn new(obj: &Bound<'_, PyAny>) -> PyResult { + pub fn new(obj: Bound<'_, PyAny>, session: Option>) -> PyResult { + let py = obj.py(); if let Ok(py_table) = obj.extract::() { Ok(py_table) } else if let Ok(py_table) = obj .getattr("_inner") - .and_then(|inner| inner.extract::()) + .and_then(|inner| inner.extract::().map_err(Into::::into)) { Ok(py_table) } else if let Ok(py_df) = obj.extract::() { @@ -73,15 +86,20 @@ impl PyTable { Ok(PyTable::from(provider)) } else if let Ok(py_df) = obj .getattr("df") - .and_then(|inner| inner.extract::()) + .and_then(|inner| inner.extract::().map_err(Into::::into)) { let provider = py_df.inner_df().as_ref().clone().into_view(); Ok(PyTable::from(provider)) - } else if let Some(provider) = table_provider_from_pycapsule(obj)? { + } else if let Some(provider) = { + let session = match session { + Some(session) => session, + None => PySessionContext::global_ctx()?.into_bound_py_any(obj.py())?, + }; + table_provider_from_pycapsule(obj.clone(), session)? + } { Ok(PyTable::from(provider)) } else { - let py = obj.py(); - let provider = Arc::new(Dataset::new(obj, py)?) as Arc; + let provider = Arc::new(Dataset::new(&obj, py)?) as Arc; Ok(PyTable::from(provider)) } } @@ -131,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()) } @@ -192,3 +206,51 @@ impl TableProvider for TempViewTable { Ok(vec![TableProviderFilterPushDown::Exact; filters.len()]) } } + +#[derive(Debug)] +pub(crate) struct RustWrappedPyTableProviderFactory { + pub(crate) table_provider_factory: Py, + pub(crate) codec: Arc, +} + +impl RustWrappedPyTableProviderFactory { + pub fn new(table_provider_factory: Py, codec: Arc) -> Self { + Self { + table_provider_factory, + codec, + } + } + + fn create_inner( + &self, + cmd: CreateExternalTable, + codec: Bound, + ) -> PyResult> { + Python::attach(|py| { + let provider = self.table_provider_factory.bind(py); + let cmd = PyCreateExternalTable::from(cmd); + + provider + .call_method1("create", (cmd,)) + .and_then(|t| PyTable::new(t, Some(codec))) + .map(|t| t.table()) + }) + } +} + +#[async_trait] +impl TableProviderFactory for RustWrappedPyTableProviderFactory { + async fn create( + &self, + _: &dyn Session, + cmd: &CreateExternalTable, + ) -> datafusion::common::Result> { + Python::attach(|py| { + let codec = create_logical_extension_capsule(py, self.codec.as_ref()) + .map_err(errors::to_datafusion_err)?; + + self.create_inner(cmd.clone(), codec.into_any()) + .map_err(errors::to_datafusion_err) + }) + } +} diff --git a/crates/core/src/udaf.rs b/crates/core/src/udaf.rs new file mode 100644 index 000000000..caf7b97bc --- /dev/null +++ b/crates/core/src/udaf.rs @@ -0,0 +1,394 @@ +// 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::ptr::NonNull; +use std::sync::Arc; + +use datafusion::arrow::array::ArrayRef; +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, Signature, Volatility, +}; +use datafusion_ffi::udaf::FFI_AggregateUDF; +use datafusion_python_util::parse_volatility; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyTuple}; + +use crate::common::data_type::PyScalarValue; +use crate::errors::{PyDataFusionResult, py_datafusion_err, to_datafusion_err}; +use crate::expr::PyExpr; + +#[derive(Debug)] +struct RustAccumulator { + accum: Py, +} + +impl RustAccumulator { + fn new(accum: Py) -> Self { + Self { accum } + } +} + +impl Accumulator for RustAccumulator { + fn state(&mut self) -> Result> { + Python::attach(|py| -> PyResult> { + let values = self.accum.bind(py).call_method0("state")?; + let mut scalars = Vec::new(); + for item in values.try_iter()? { + let item: Bound<'_, PyAny> = item?; + let scalar = item.extract::()?.0; + scalars.push(scalar); + } + Ok(scalars) + }) + .map_err(|e| DataFusionError::Execution(format!("{e}"))) + } + + fn evaluate(&mut self) -> Result { + Python::attach(|py| -> PyResult { + let value = self.accum.bind(py).call_method0("evaluate")?; + value.extract::().map(|v| v.0) + }) + .map_err(|e| DataFusionError::Execution(format!("{e}"))) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + Python::attach(|py| { + // 1. cast args to Pyarrow array + let py_args = values + .iter() + .map(|arg| arg.to_data().to_pyarrow(py).unwrap()) + .collect::>(); + let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?; + + // 2. call function + self.accum + .bind(py) + .call_method1("update", py_args) + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + + Ok(()) + }) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + Python::attach(|py| { + // // 1. cast states to Pyarrow arrays + let py_states: Result>> = states + .iter() + .map(|state| { + state + .to_data() + .to_pyarrow(py) + .map_err(|e| DataFusionError::Execution(format!("{e}"))) + }) + .collect(); + + // 2. call merge + self.accum + .bind(py) + .call_method1("merge", (py_states?,)) + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + + Ok(()) + }) + } + + fn size(&self) -> usize { + std::mem::size_of_val(self) + } + + fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + Python::attach(|py| { + // 1. cast args to Pyarrow array + let py_args = values + .iter() + .map(|arg| arg.to_data().to_pyarrow(py).unwrap()) + .collect::>(); + let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?; + + // 2. call function + self.accum + .bind(py) + .call_method1("retract_batch", py_args) + .map_err(|e| DataFusionError::Execution(format!("{e}")))?; + + Ok(()) + }) + } + + fn supports_retract_batch(&self) -> bool { + Python::attach( + |py| match self.accum.bind(py).call_method0("supports_retract_batch") { + Ok(x) => x.extract().unwrap_or(false), + Err(_) => false, + }, + ) + } +} + +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| 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 { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_aggregate_udf"))? + .cast(); + let udaf = unsafe { data.as_ref() }; + let udaf: Arc = udaf.into(); + + Ok(AggregateUDF::new_from_shared_impl(udaf)) +} + +/// Represents an AggregateUDF +#[pyclass( + from_py_object, + frozen, + name = "AggregateUDF", + module = "datafusion", + subclass +)] +#[derive(Debug, Clone)] +pub struct PyAggregateUDF { + pub(crate) function: AggregateUDF, +} + +#[pymethods] +impl PyAggregateUDF { + #[new] + #[pyo3(signature=(name, accumulator, input_type, return_type, state_type, volatility))] + fn new( + name: &str, + accumulator: Py, + input_type: PyArrowType>, + return_type: PyArrowType, + state_type: PyArrowType>, + volatility: &str, + ) -> PyResult { + let py_udf = PythonFunctionAggregateUDF::new( + name.to_string(), + accumulator, + input_type.0, + return_type.0, + state_type.0, + parse_volatility(volatility)?, + ); + let function = AggregateUDF::new_from_impl(py_udf); + Ok(Self { function }) + } + + #[staticmethod] + pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { + if func.is_instance_of::() { + let capsule = func.cast::().map_err(py_datafusion_err)?; + let function = aggregate_udf_from_capsule(capsule)?; + return Ok(Self { function }); + } + + if func.hasattr("__datafusion_aggregate_udf__")? { + let capsule = func.getattr("__datafusion_aggregate_udf__")?.call0()?; + let capsule = capsule.cast::().map_err(py_datafusion_err)?; + let function = aggregate_udf_from_capsule(capsule)?; + return Ok(Self { function }); + } + + Err(crate::errors::PyDataFusionError::Common( + "__datafusion_aggregate_udf__ does not exist on AggregateUDF object.".to_string(), + )) + } + + /// creates a new PyExpr with the call of the udf + #[pyo3(signature = (*args))] + fn __call__(&self, args: Vec) -> PyResult { + let args = args.iter().map(|e| e.expr.clone()).collect(); + Ok(self.function.call(args).into()) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("AggregateUDF({})", self.function.name())) + } + + #[getter] + fn name(&self) -> &str { + self.function.name() + } +} diff --git a/crates/core/src/udf.rs b/crates/core/src/udf.rs new file mode 100644 index 000000000..2006401db --- /dev/null +++ b/crates/core/src/udf.rs @@ -0,0 +1,284 @@ +// 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::hash::{Hash, Hasher}; +use std::ptr::NonNull; +use std::sync::Arc; + +use arrow::datatypes::{Field, FieldRef}; +use arrow::pyarrow::ToPyArrow; +use datafusion::arrow::array::{ArrayData, make_array}; +use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType}; +use datafusion::common::internal_err; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_ffi::udf::FFI_ScalarUDF; +use datafusion_python_util::parse_volatility; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyTuple}; + +use crate::array::PyArrowArrayExportable; +use crate::errors::{PyDataFusionResult, to_datafusion_err}; +use crate::expr::PyExpr; + +/// This struct holds the Python written function that is a +/// ScalarUDF. +#[derive(Debug)] +pub(crate) struct PythonFunctionScalarUDF { + name: String, + func: Py, + signature: Signature, + return_field: FieldRef, +} + +impl PythonFunctionScalarUDF { + fn new( + name: String, + func: Py, + input_fields: Vec, + return_field: Field, + volatility: Volatility, + ) -> Self { + let input_types = input_fields.iter().map(|f| f.data_type().clone()).collect(); + let signature = Signature::exact(input_types, volatility); + Self { + name, + func, + signature, + 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 {} +impl PartialEq for PythonFunctionScalarUDF { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.signature == other.signature + && self.return_field == other.return_field + // 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); + } +} + +impl ScalarUDFImpl for PythonFunctionScalarUDF { + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> datafusion::common::Result { + internal_err!( + "return_field should not be called when return_field_from_args is implemented." + ) + } + + fn return_field_from_args( + &self, + _args: ReturnFieldArgs, + ) -> datafusion::common::Result { + Ok(Arc::clone(&self.return_field)) + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion::common::Result { + let num_rows = args.number_rows; + Python::attach(|py| { + // 1. cast args to Pyarrow arrays + let py_args = args + .args + .into_iter() + .zip(args.arg_fields) + .map(|(arg, field)| { + let array = arg.to_array(num_rows)?; + PyArrowArrayExportable::new(array, field) + .to_pyarrow(py) + .map_err(to_datafusion_err) + }) + .collect::, _>>()?; + let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?; + + // 2. call function + let value = self + .func + .call(py, py_args, None) + .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + + // 3. cast to arrow::array::Array + let array_data = ArrayData::from_pyarrow_bound(value.bind(py)) + .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + Ok(ColumnarValue::Array(make_array(array_data))) + }) + } +} + +/// Represents a PyScalarUDF +#[pyclass( + from_py_object, + frozen, + name = "ScalarUDF", + module = "datafusion", + subclass +)] +#[derive(Debug, Clone)] +pub struct PyScalarUDF { + pub(crate) function: ScalarUDF, +} + +#[pymethods] +impl PyScalarUDF { + #[new] + #[pyo3(signature=(name, func, input_types, return_type, volatility))] + fn new( + name: String, + func: Py, + input_types: PyArrowType>, + return_type: PyArrowType, + volatility: &str, + ) -> PyResult { + let py_function = PythonFunctionScalarUDF::new( + name, + func, + input_types.0, + return_type.0, + parse_volatility(volatility)?, + ); + let function = ScalarUDF::new_from_impl(py_function); + + Ok(Self { function }) + } + + #[staticmethod] + pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { + if func.hasattr("__datafusion_scalar_udf__")? { + let capsule = func.getattr("__datafusion_scalar_udf__")?.call0()?; + let capsule = capsule.cast::().map_err(to_datafusion_err)?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_scalar_udf"))? + .cast(); + let udf = unsafe { data.as_ref() }; + let udf: Arc = udf.into(); + + Ok(Self { + function: ScalarUDF::new_from_shared_impl(udf), + }) + } else { + Err(crate::errors::PyDataFusionError::Common( + "__datafusion_scalar_udf__ does not exist on ScalarUDF object.".to_string(), + )) + } + } + + /// creates a new PyExpr with the call of the udf + #[pyo3(signature = (*args))] + fn __call__(&self, args: Vec) -> PyResult { + let args = args.iter().map(|e| e.expr.clone()).collect(); + Ok(self.function.call(args).into()) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("ScalarUDF({})", self.function.name())) + } + + #[getter] + fn name(&self) -> &str { + self.function.name() + } +} diff --git a/crates/core/src/udtf.rs b/crates/core/src/udtf.rs new file mode 100644 index 000000000..cffa0c12a --- /dev/null +++ b/crates/core/src/udtf.rs @@ -0,0 +1,202 @@ +// 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::ptr::NonNull; +use std::sync::Arc; + +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, 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)] +pub struct PyTableFunction { + pub(crate) name: String, + pub(crate) inner: PyTableFunctionInner, +} + +#[derive(Debug, Clone)] +pub(crate) enum PyTableFunctionInner { + PythonFunction(PythonTableFunctionCallable), + FFIFunction(Arc), +} + +#[pymethods] +impl PyTableFunction { + #[new] + #[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(); + let session = match session { + Some(session) => session, + None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, + }; + let capsule = func + .getattr("__datafusion_table_function__")? + .call1((session,)).map_err(|err| { + if err.get_type(py).is(PyType::new::(py)) { + PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table functions. Either downgrade DataFusion or upgrade your function library.") + } else { + err + } + })?; + let capsule = capsule.cast::()?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_table_function"))? + .cast(); + let ffi_func = unsafe { data.as_ref() }; + let foreign_func: Arc = ffi_func.to_owned().into(); + + PyTableFunctionInner::FFIFunction(foreign_func) + } else { + PyTableFunctionInner::PythonFunction(PythonTableFunctionCallable { + callable: Arc::new(func.unbind()), + inject_session_on_call, + }) + }; + + Ok(Self { + name: name.to_string(), + inner, + }) + } + + #[pyo3(signature = (*args))] + pub fn __call__(&self, args: Vec) -> PyResult { + let args: Vec = args.iter().map(|e| e.expr.clone()).collect(); + 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)) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("TableUDF({})", self.name)) + } +} + +/// 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: &PythonTableFunctionCallable, + args: TableFunctionArgs, +) -> DataFusionResult> { + 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::>(); + + Python::attach(|py| { + 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) + }) + .map_err(to_datafusion_err) +} + +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(callable) => { + call_python_table_function(callable, args) + } + } + } +} diff --git a/src/udwf.rs b/crates/core/src/udwf.rs similarity index 64% rename from src/udwf.rs rename to crates/core/src/udwf.rs index d347ec0f1..ebec8f3bd 100644 --- a/src/udwf.rs +++ b/crates/core/src/udwf.rs @@ -15,30 +15,29 @@ // 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; -use arrow::array::{make_array, Array, ArrayData, ArrayRef}; +use arrow::array::{Array, ArrayData, ArrayRef, make_array}; 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, }; use datafusion::scalar::ScalarValue; -use datafusion_ffi::udwf::{FFI_WindowUDF, ForeignWindowUDF}; +use datafusion_ffi::udwf::FFI_WindowUDF; +use datafusion_python_util::parse_volatility; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyList, PyTuple}; use crate::common::data_type::PyScalarValue; -use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult}; +use crate::errors::{PyDataFusionResult, to_datafusion_err}; use crate::expr::PyExpr; -use crate::utils::{parse_volatility, validate_pycapsule}; #[derive(Debug)] struct RustPartitionEvaluator { @@ -94,7 +93,6 @@ impl PartitionEvaluator for RustPartitionEvaluator { } fn evaluate_all(&mut self, values: &[ArrayRef], num_rows: usize) -> Result { - println!("evaluate all called with number of values {}", values.len()); Python::attach(|py| { let py_values = PyList::new( py, @@ -198,19 +196,34 @@ 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 -#[pyclass(frozen, name = "WindowUDF", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "WindowUDF", + module = "datafusion", + subclass +)] #[derive(Debug, Clone)] pub struct PyWindowUDF { pub(crate) function: WindowUDF, @@ -228,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 }) } @@ -249,72 +262,125 @@ impl PyWindowUDF { #[staticmethod] pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { - if func.hasattr("__datafusion_window_udf__")? { - let capsule = func.getattr("__datafusion_window_udf__")?.call0()?; - let capsule = capsule.downcast::().map_err(py_datafusion_err)?; - validate_pycapsule(capsule, "datafusion_window_udf")?; - - let udwf = unsafe { capsule.reference::() }; - let udwf: ForeignWindowUDF = udwf.try_into()?; - - Ok(Self { - function: udwf.into(), - }) + let capsule = if func.hasattr("__datafusion_window_udf__")? { + func.getattr("__datafusion_window_udf__")?.call0()? } else { - Err(crate::errors::PyDataFusionError::Common( - "__datafusion_window_udf__ does not exist on WindowUDF object.".to_string(), - )) - } + func + }; + + let capsule = capsule.cast::().map_err(to_datafusion_err)?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_window_udf"))? + .cast(); + let udwf = unsafe { data.as_ref() }; + let udwf: Arc = udwf.into(); + + Ok(Self { + function: WindowUDF::new_from_shared_impl(udwf), + }) } 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 WindowUDFImpl for MultiColumnWindowUDF { - fn as_any(&self) -> &dyn Any { - self +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 PythonFunctionWindowUDF { fn name(&self) -> &str { &self.name } @@ -333,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/src/unparser/dialect.rs b/crates/core/src/unparser/dialect.rs similarity index 93% rename from src/unparser/dialect.rs rename to crates/core/src/unparser/dialect.rs index 5df0a0c2e..52a2da00b 100644 --- a/src/unparser/dialect.rs +++ b/crates/core/src/unparser/dialect.rs @@ -22,7 +22,13 @@ use datafusion::sql::unparser::dialect::{ }; use pyo3::prelude::*; -#[pyclass(frozen, name = "Dialect", module = "datafusion.unparser", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Dialect", + module = "datafusion.unparser", + subclass +)] #[derive(Clone)] pub struct PyDialect { pub dialect: Arc, diff --git a/src/unparser/mod.rs b/crates/core/src/unparser/mod.rs similarity index 94% rename from src/unparser/mod.rs rename to crates/core/src/unparser/mod.rs index 908b59d3b..5142b918e 100644 --- a/src/unparser/mod.rs +++ b/crates/core/src/unparser/mod.rs @@ -19,15 +19,21 @@ mod dialect; use std::sync::Arc; -use datafusion::sql::unparser::dialect::Dialect; use datafusion::sql::unparser::Unparser; +use datafusion::sql::unparser::dialect::Dialect; use dialect::PyDialect; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use crate::sql::logical::PyLogicalPlan; -#[pyclass(frozen, name = "Unparser", module = "datafusion.unparser", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Unparser", + module = "datafusion.unparser", + subclass +)] #[derive(Clone)] pub struct PyUnparser { dialect: Arc, diff --git a/crates/util/Cargo.toml b/crates/util/Cargo.toml new file mode 100644 index 000000000..c23667b0f --- /dev/null +++ b/crates/util/Cargo.toml @@ -0,0 +1,35 @@ +# 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. + +[package] +name = "datafusion-python-util" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description.workspace = true +homepage.workspace = true +repository.workspace = true + +[dependencies] +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/src/errors.rs b/crates/util/src/errors.rs similarity index 87% rename from src/errors.rs rename to crates/util/src/errors.rs index fc079eb6c..0d25c8847 100644 --- a/src/errors.rs +++ b/crates/util/src/errors.rs @@ -22,8 +22,8 @@ use std::fmt::Debug; use datafusion::arrow::error::ArrowError; use datafusion::error::DataFusionError as InnerDataFusionError; use prost::EncodeError; -use pyo3::exceptions::PyException; use pyo3::PyErr; +use pyo3::exceptions::{PyException, PyValueError}; pub type PyDataFusionResult = std::result::Result; @@ -39,7 +39,7 @@ pub enum PyDataFusionError { impl fmt::Display for PyDataFusionError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { - PyDataFusionError::ExecutionError(e) => write!(f, "DataFusion error: {e:?}"), + PyDataFusionError::ExecutionError(e) => write!(f, "DataFusion error: {e}"), PyDataFusionError::ArrowError(e) => write!(f, "Arrow error: {e:?}"), PyDataFusionError::PythonError(e) => write!(f, "Python error {e:?}"), PyDataFusionError::Common(e) => write!(f, "{e}"), @@ -96,3 +96,13 @@ pub fn py_unsupported_variant_err(e: impl Debug) -> PyErr { pub fn to_datafusion_err(e: impl Debug) -> InnerDataFusionError { InnerDataFusionError::Execution(format!("{e:?}")) } + +pub fn from_datafusion_error(err: InnerDataFusionError) -> PyErr { + match err { + InnerDataFusionError::External(boxed) => match boxed.downcast::() { + Ok(py_err) => *py_err, + Err(original_boxed) => PyValueError::new_err(format!("{original_boxed}")), + }, + _ => PyValueError::new_err(format!("{err}")), + } +} diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs new file mode 100644 index 000000000..9327d7f2f --- /dev/null +++ b/crates/util/src/lib.rs @@ -0,0 +1,348 @@ +// 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::future::Future; +use std::ptr::NonNull; +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::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; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyType}; +use tokio::runtime::Runtime; +use tokio::task::JoinHandle; +use tokio::time::sleep; + +pub mod errors; +pub use crate::errors::to_datafusion_err; +use crate::errors::{PyDataFusionError, PyDataFusionResult}; + +/// Utility to get the Tokio Runtime from Python +#[inline] +pub fn get_tokio_runtime() -> &'static Runtime { + // NOTE: Other pyo3 python libraries have had issues with using tokio + // behind a forking app-server like `gunicorn` + // If we run into that problem, in the future we can look to `delta-rs` + // which adds a check in that disallows calls from a forked process + // https://github.com/delta-io/delta-rs/blob/87010461cfe01563d91a4b9cd6fa468e2ad5f283/python/src/utils.rs#L10-L31 + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| Runtime::new().unwrap()) +} + +#[inline] +pub fn is_ipython_env(py: Python) -> &'static bool { + static IS_IPYTHON_ENV: OnceLock = OnceLock::new(); + IS_IPYTHON_ENV.get_or_init(|| { + py.import("IPython") + .and_then(|ipython| ipython.call_method0("get_ipython")) + .map(|ipython| !ipython.is_none()) + .unwrap_or(false) + }) +} + +/// Utility to get the Global Datafussion CTX +#[inline] +pub fn get_global_ctx() -> &'static Arc { + static CTX: OnceLock> = OnceLock::new(); + CTX.get_or_init(|| Arc::new(SessionContext::new())) +} + +/// Utility to collect rust futures with GIL released and respond to +/// Python interrupts such as ``KeyboardInterrupt``. If a signal is +/// received while the future is running, the future is aborted and the +/// corresponding Python exception is raised. +pub fn wait_for_future(py: Python, fut: F) -> PyResult +where + F: Future + Send, + F::Output: Send, +{ + let runtime: &Runtime = get_tokio_runtime(); + const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(1_000); + + // Some fast running processes that generate many `wait_for_future` calls like + // PartitionedDataFrameStreamReader::next require checking for interrupts early + py.run(cr"pass", None, None)?; + py.check_signals()?; + + py.detach(|| { + runtime.block_on(async { + tokio::pin!(fut); + loop { + tokio::select! { + res = &mut fut => break Ok(res), + _ = sleep(INTERVAL_CHECK_SIGNALS) => { + Python::attach(|py| { + // Execute a no-op Python statement to trigger signal processing. + // This is necessary because py.check_signals() alone doesn't + // actually check for signals - it only raises an exception if + // a signal was already set during a previous Python API call. + // Running even trivial Python code forces the interpreter to + // process any pending signals (like KeyboardInterrupt). + py.run(cr"pass", None, None)?; + py.check_signals() + })?; + } + } + } + }) + }) +} + +/// Spawn a [`Future`] on the Tokio runtime and wait for completion +/// while respecting Python signal handling. +pub fn spawn_future(py: Python, fut: F) -> PyDataFusionResult +where + F: Future> + Send + 'static, + T: Send + 'static, +{ + let rt = get_tokio_runtime(); + let handle: JoinHandle> = rt.spawn(fut); + // Wait for the join handle while respecting Python signal handling. + // We handle errors in two steps so `?` maps the error types correctly: + // 1) convert any Python-related error from `wait_for_future` into `PyDataFusionError` + // 2) convert any DataFusion error (inner result) into `PyDataFusionError` + let inner_result = wait_for_future(py, async { + // handle.await yields `Result, JoinError>` + // map JoinError into a DataFusion error so the async block returns + // `datafusion::common::Result` (i.e. Result) + match handle.await { + Ok(inner) => inner, + Err(join_err) => Err(to_datafusion_err(join_err)), + } + })?; // converts PyErr -> PyDataFusionError + + // `inner_result` is `datafusion::common::Result`; use `?` to convert + // the inner DataFusion error into `PyDataFusionError` via `From` and + // return the inner `T` on success. + Ok(inner_result?) +} + +pub fn parse_volatility(value: &str) -> PyDataFusionResult { + Ok(match value { + "immutable" => Volatility::Immutable, + "stable" => Volatility::Stable, + "volatile" => Volatility::Volatile, + value => { + return Err(PyDataFusionError::Common(format!( + "Unsupported volatility type: `{value}`, supported \ + values are: immutable, stable and volatile." + ))); + } + }) +} + +pub fn validate_pycapsule(capsule: &Bound, name: &str) -> PyResult<()> { + let capsule_name = capsule.name()?; + if capsule_name.is_none() { + return Err(PyValueError::new_err(format!( + "Expected {name} PyCapsule to have name set." + ))); + } + + let capsule_name = unsafe { capsule_name.unwrap().as_cstr().to_str() } + .map_err(|err| PyValueError::new_err(err.to_string()))?; + if capsule_name != name { + return Err(PyValueError::new_err(format!( + "Expected name '{name}' in PyCapsule, instead got '{capsule_name}'" + ))); + } + + Ok(()) +} + +pub fn table_provider_from_pycapsule<'py>( + mut obj: Bound<'py, PyAny>, + session: Bound<'py, PyAny>, +) -> PyResult>> { + if obj.hasattr("__datafusion_table_provider__")? { + obj = obj + .getattr("__datafusion_table_provider__")? + .call1((session,)).map_err(|err| { + let py = obj.py(); + if err.get_type(py).is(PyType::new::(py)) { + PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table providers. Either downgrade DataFusion or upgrade your function library.") + } else { + err + } + })?; + } + + if let Ok(capsule) = obj.cast::() { + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_table_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + + Ok(Some(provider)) + } else { + Ok(None) + } +} + +pub fn create_logical_extension_capsule<'py>( + py: Python<'py>, + codec: &FFI_LogicalExtensionCodec, +) -> PyResult> { + let codec = codec.clone(); + + PyCapsule::new_with_value(py, codec, cr"datafusion_logical_extension_codec") +} + +pub fn ffi_logical_codec_from_pycapsule(obj: Bound) -> PyResult { + let attr_name = "__datafusion_logical_extension_codec__"; + let capsule = if obj.hasattr(attr_name)? { + obj.getattr(attr_name)?.call0()? + } else { + obj + }; + + let capsule = capsule.cast::()?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_logical_extension_codec"))? + .cast(); + let codec = unsafe { data.as_ref() }; + + Ok(codec.clone()) +} + +pub fn create_physical_extension_capsule<'py>( + py: Python<'py>, + codec: &FFI_PhysicalExtensionCodec, +) -> PyResult> { + let codec = codec.clone(); + + PyCapsule::new_with_value(py, codec, cr"datafusion_physical_extension_codec") +} + +/// 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 +); + +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", + FFI_TaskContextProvider, + TaskContext +); diff --git a/dev/changelog/52.0.0.md b/dev/changelog/52.0.0.md new file mode 100644 index 000000000..3f848bb47 --- /dev/null +++ b/dev/changelog/52.0.0.md @@ -0,0 +1,78 @@ + + +# Apache DataFusion Python 52.0.0 Changelog + +This release consists of 26 commits from 9 contributors. See credits at the end of this changelog for more information. + +**Implemented enhancements:** + +- feat: add CatalogProviderList support [#1363](https://github.com/apache/datafusion-python/pull/1363) (timsaucer) +- feat: add support for generating JSON formatted substrait plan [#1376](https://github.com/apache/datafusion-python/pull/1376) (Prathamesh9284) +- feat: add regexp_instr function [#1382](https://github.com/apache/datafusion-python/pull/1382) (mesejo) + +**Fixed bugs:** + +- fix: mangled errors [#1377](https://github.com/apache/datafusion-python/pull/1377) (mesejo) + +**Documentation updates:** + +- docs: Clarify first_value usage in select vs aggregate [#1348](https://github.com/apache/datafusion-python/pull/1348) (AdMub) + +**Other:** + +- Release 51.0.0 [#1333](https://github.com/apache/datafusion-python/pull/1333) (timsaucer) +- Use explicit timer in unit test [#1338](https://github.com/apache/datafusion-python/pull/1338) (timsaucer) +- Add use_fabric_endpoint parameter to MicrosoftAzure class [#1357](https://github.com/apache/datafusion-python/pull/1357) (djouallah) +- Prepare for DF52 release [#1337](https://github.com/apache/datafusion-python/pull/1337) (timsaucer) +- build(deps): bump actions/checkout from 5 to 6 [#1310](https://github.com/apache/datafusion-python/pull/1310) (dependabot[bot]) +- build(deps): bump actions/download-artifact from 5 to 7 [#1321](https://github.com/apache/datafusion-python/pull/1321) (dependabot[bot]) +- build(deps): bump actions/upload-artifact from 4 to 6 [#1322](https://github.com/apache/datafusion-python/pull/1322) (dependabot[bot]) +- build(deps): bump actions/cache from 4 to 5 [#1323](https://github.com/apache/datafusion-python/pull/1323) (dependabot[bot]) +- Pass Field information back and forth when using scalar UDFs [#1299](https://github.com/apache/datafusion-python/pull/1299) (timsaucer) +- Update dependency minor versions to prepare for DF52 release [#1368](https://github.com/apache/datafusion-python/pull/1368) (timsaucer) +- Improve displayed error by using `DataFusionError`'s `Display` trait [#1370](https://github.com/apache/datafusion-python/pull/1370) (abey79) +- Enforce DataFrame display memory limits with `max_rows` + `min_rows` constraint (deprecate `repr_rows`) [#1367](https://github.com/apache/datafusion-python/pull/1367) (kosiew) +- Implement all CSV reader options [#1361](https://github.com/apache/datafusion-python/pull/1361) (timsaucer) +- chore: add confirmation before tarball is released [#1372](https://github.com/apache/datafusion-python/pull/1372) (milenkovicm) +- Build in debug mode for PRs [#1375](https://github.com/apache/datafusion-python/pull/1375) (timsaucer) +- minor: remove ffi test wheel from distribution artifact [#1378](https://github.com/apache/datafusion-python/pull/1378) (timsaucer) +- chore: update rust 2024 edition [#1371](https://github.com/apache/datafusion-python/pull/1371) (timsaucer) +- Fix Python UDAF list-of-timestamps return by enforcing list-valued scalars and caching PyArrow types [#1347](https://github.com/apache/datafusion-python/pull/1347) (kosiew) +- minor: update cargo dependencies [#1383](https://github.com/apache/datafusion-python/pull/1383) (timsaucer) +- chore: bump Python version for RAT checking [#1386](https://github.com/apache/datafusion-python/pull/1386) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 13 Tim Saucer + 4 dependabot[bot] + 2 Daniel Mesejo + 2 kosiew + 1 Adisa Mubarak (AdMub) + 1 Antoine Beyeler + 1 Dhanashri Prathamesh Iranna + 1 Marko Milenković + 1 Mimoune +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + 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. + diff --git a/dev/changelog/54.0.0.md b/dev/changelog/54.0.0.md new file mode 100644 index 000000000..e73118d8b --- /dev/null +++ b/dev/changelog/54.0.0.md @@ -0,0 +1,106 @@ + + +# Apache DataFusion Python 54.0.0 Changelog + +This release consists of 51 commits from 8 contributors. See credits at the end of this changelog for more information. + +**Breaking changes:** + +- Update datafusion dependency to latest in preparation for DF54 [#1532](https://github.com/apache/datafusion-python/pull/1532) (timsaucer) +- feat: enable pickling for Python aggregate and window UDFs [#1545](https://github.com/apache/datafusion-python/pull/1545) (timsaucer) +- feat: pass calling SessionContext to Python UDTF callbacks [#1555](https://github.com/apache/datafusion-python/pull/1555) (timsaucer) +- feat: accept distinct kwarg on sum and avg [#1556](https://github.com/apache/datafusion-python/pull/1556) (timsaucer) + +**Implemented enhancements:** + +- feat: add AI skill to find and improve the Pythonic interface to functions [#1484](https://github.com/apache/datafusion-python/pull/1484) (timsaucer) +- feat: enable pickling of most Expr except udaf and udwf [#1544](https://github.com/apache/datafusion-python/pull/1544) (timsaucer) +- feat: expose variety of features from DF54 update [#1554](https://github.com/apache/datafusion-python/pull/1554) (timsaucer) +- feat: Python UDFs: per-session inlining toggle and strict refusal setting [#1546](https://github.com/apache/datafusion-python/pull/1546) (timsaucer) +- feat: create free-threaded python wheels [#1553](https://github.com/apache/datafusion-python/pull/1553) (timsaucer) +- feat: expose lambda and higher-order array functions [#1561](https://github.com/apache/datafusion-python/pull/1561) (timsaucer) +- feat: import user-defined physical optimizer rules over FFI [#1557](https://github.com/apache/datafusion-python/pull/1557) (timsaucer) +- feat: expose SessionContext.copied_config and parse_capacity_limit [#1570](https://github.com/apache/datafusion-python/pull/1570) (timsaucer) +- feat: expose array_compact, array_normalize, cosine_distance, inner_product [#1567](https://github.com/apache/datafusion-python/pull/1567) (timsaucer) +- feat: expose arrow_field, arrow_try_cast, cast_to_type, with_metadata [#1568](https://github.com/apache/datafusion-python/pull/1568) (timsaucer) +- feat: expose spark-compatible functions [#1564](https://github.com/apache/datafusion-python/pull/1564) (timsaucer) +- feat: improve pythonic interface on date/time functions [#1563](https://github.com/apache/datafusion-python/pull/1563) (timsaucer) + +**Fixed bugs:** + +- fix: type scalar UDF returns as Arrow arrays [#1528](https://github.com/apache/datafusion-python/pull/1528) (BharatDeva) +- fix: Skip `fork` and `forkserver` on `win32` [#1566](https://github.com/apache/datafusion-python/pull/1566) (nuno-faria) + +**Documentation updates:** + +- docs: enrich module docstrings and add doctest examples [#1498](https://github.com/apache/datafusion-python/pull/1498) (timsaucer) +- docs: add README section for AI coding assistants [#1503](https://github.com/apache/datafusion-python/pull/1503) (timsaucer) +- docs: add upstream sync process documentation [#1524](https://github.com/apache/datafusion-python/pull/1524) (timsaucer) +- docs: document null-handling function arguments [#1527](https://github.com/apache/datafusion-python/pull/1527) (BharatDeva) +- docs: user guide + runnable examples for distributing expressions [#1547](https://github.com/apache/datafusion-python/pull/1547) (timsaucer) +- docs: convert reStructuredText sources to MyST markdown [#1579](https://github.com/apache/datafusion-python/pull/1579) (timsaucer) + +**Other:** + +- Release 53.0.0 [#1491](https://github.com/apache/datafusion-python/pull/1491) (timsaucer) +- ci: disable symbol export on Windows verification [#1486](https://github.com/apache/datafusion-python/pull/1486) (timsaucer) +- Add Python bindings for accessing ExecutionMetrics [#1381](https://github.com/apache/datafusion-python/pull/1381) (ShreyeshArangath) +- Support None comparisons for null expressions [#1489](https://github.com/apache/datafusion-python/pull/1489) (zeel2104) +- chore: update release documentation [#1494](https://github.com/apache/datafusion-python/pull/1494) (timsaucer) +- Fix error on show() with an explain plan [#1492](https://github.com/apache/datafusion-python/pull/1492) (timsaucer) +- Add SKILL.md and enrich package docstring [#1497](https://github.com/apache/datafusion-python/pull/1497) (timsaucer) +- minor: fix header on agent SKILL.md file [#1501](https://github.com/apache/datafusion-python/pull/1501) (timsaucer) +- tpch examples: rewrite queries idiomatically and embed reference SQL [#1504](https://github.com/apache/datafusion-python/pull/1504) (timsaucer) +- Move public skills to a directory to avoid downloading the whole repo [#1519](https://github.com/apache/datafusion-python/pull/1519) (ntjohnson1) +- Update user documentation for AI agent skill usage [#1505](https://github.com/apache/datafusion-python/pull/1505) (timsaucer) +- build(deps): combined dependabot bumps (Cargo + workflows) [#1534](https://github.com/apache/datafusion-python/pull/1534) (timsaucer) +- Add support for logical and physical codecs [#1541](https://github.com/apache/datafusion-python/pull/1541) (timsaucer) +- Add details on caching to skill [#1521](https://github.com/apache/datafusion-python/pull/1521) (ntjohnson1) +- Bump DataFusion to prepare for DF54 release candidate [#1562](https://github.com/apache/datafusion-python/pull/1562) (timsaucer) +- Export `to_datafusion_err` from the util crate root [#1487](https://github.com/apache/datafusion-python/pull/1487) (kosiew) +- chore: remove unused PyConfig [#1485](https://github.com/apache/datafusion-python/pull/1485) (timsaucer) +- refactor(context): deduplicate register/read option-building logic [#1479](https://github.com/apache/datafusion-python/pull/1479) (mesejo) +- Improve documentation site layout [#1578](https://github.com/apache/datafusion-python/pull/1578) (timsaucer) +- Allow `DataFrame.aggregate` to accept `None` for no grouping [#1581](https://github.com/apache/datafusion-python/pull/1581) (kosiew) +- Update to released DF 54.0.0 [#1588](https://github.com/apache/datafusion-python/pull/1588) (timsaucer) +- build(deps): batch dependabot dependency updates [#1589](https://github.com/apache/datafusion-python/pull/1589) (timsaucer) +- Deprecate `Expr` temporal part arguments in date extraction and truncation functions [#1587](https://github.com/apache/datafusion-python/pull/1587) (kosiew) +- chore: update rust dependencies [#1604](https://github.com/apache/datafusion-python/pull/1604) (timsaucer) +- Add deprecation warnings for Expr passed to confirmed literal-only function arguments [#1605](https://github.com/apache/datafusion-python/pull/1605) (kosiew) +- chore: resolve audits after DF54.0.0 update to skill [#1608](https://github.com/apache/datafusion-python/pull/1608) (timsaucer) +- chore: remove 3.13 freethreaded builds [#1609](https://github.com/apache/datafusion-python/pull/1609) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 39 Tim Saucer + 4 kosiew + 2 BharatDeva + 2 Nick + 1 Daniel Mesejo + 1 Nuno Faria + 1 Shreyesh + 1 Zeel Desai +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/check_crates_patch.py b/dev/check_crates_patch.py new file mode 100644 index 000000000..74e489e1f --- /dev/null +++ b/dev/check_crates_patch.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +# 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. + +"""Check that no Cargo.toml files contain [patch.crates-io] entries. + +Release builds must not depend on patched crates. During development it is +common to temporarily patch crates-io dependencies, but those patches must +be removed before creating a release. + +An empty [patch.crates-io] section is allowed. +""" + +import sys +from pathlib import Path + +import tomllib + + +def main() -> int: + errors: list[str] = [] + for cargo_toml in sorted(Path().rglob("Cargo.toml")): + if "target" in cargo_toml.parts: + continue + with Path.open(cargo_toml, "rb") as f: + data = tomllib.load(f) + patch = data.get("patch", {}).get("crates-io", {}) + if patch: + errors.append(str(cargo_toml)) + for name, spec in patch.items(): + errors.append(f" {name} = {spec}") + + if errors: + print("ERROR: Release builds must not contain [patch.crates-io] entries.") + print() + for line in errors: + print(line) + print() + print("Remove all [patch.crates-io] entries before creating a release.") + return 1 + + print("OK: No [patch.crates-io] entries found.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/dev/create_license.py b/dev/create_license.py index a28a0abec..acbf8587c 100644 --- a/dev/create_license.py +++ b/dev/create_license.py @@ -22,11 +22,9 @@ import subprocess from pathlib import Path -subprocess.check_output(["cargo", "install", "cargo-license"]) data = subprocess.check_output( [ - "cargo", - "license", + "cargo-license", "--avoid-build-deps", "--avoid-dev-deps", "--do-not-bundle", diff --git a/dev/release/README.md b/dev/release/README.md index 5d2fae5a7..384ef9210 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -26,13 +26,19 @@ 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. +## 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,8 +59,12 @@ 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-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 +75,15 @@ 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 +$ uv sync --group release +$ 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 +92,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 +106,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 +136,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 +147,42 @@ 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 + +Before sending the vote email, run the manually triggered GitHub Actions workflow +"Verify Release Candidate" and confirm all matrix jobs pass across the OS/architecture matrix +(for example, Linux, macOS, and Windows runners): + +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, `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 +matrix entries and that all jobs passed. + +```text +Verification note: The manually triggered "Verify Release Candidate" workflow was run for version and rc_number across all configured OS/architecture matrix entries, and all matrix jobs completed successfully. ``` ### Send the Email @@ -164,7 +195,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: @@ -176,7 +207,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 @@ -184,7 +215,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 @@ -205,7 +236,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 @@ -213,7 +244,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 @@ -233,15 +264,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)). @@ -264,7 +295,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 @@ -280,5 +311,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 ``` diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index dcd5d9aac..70f60168f 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -47,4 +47,9 @@ benchmarks/tpch/queries/q*.sql benchmarks/tpch/create_tables.sql .cargo/config.toml **/.cargo/config.toml -uv.lock \ No newline at end of file +uv.lock +examples/tpch/answers_sf1/*.tbl +**/SKILL.md +docs/source/llms.txt +CLAUDE.md +.claude/skills \ No newline at end of file diff --git a/dev/release/release-tarball.sh b/dev/release/release-tarball.sh index 8c305a676..2b82d1bac 100755 --- a/dev/release/release-tarball.sh +++ b/dev/release/release-tarball.sh @@ -43,6 +43,13 @@ fi version=$1 rc=$2 +read -r -p "Proceed to release tarball for ${version}-rc${rc}? [y/N]: " answer +answer=${answer:-no} +if [ "${answer}" != "y" ]; then + echo "Cancelled tarball release!" + exit 1 +fi + tmp_dir=tmp-apache-datafusion-python-dist echo "Recreate temporary directory: ${tmp_dir}" 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. diff --git a/dev/release/verify-release-candidate.sh b/dev/release/verify-release-candidate.sh index 2bfce0e2d..42e3970fb 100755 --- a/dev/release/verify-release-candidate.sh +++ b/dev/release/verify-release-candidate.sh @@ -112,8 +112,17 @@ test_source_distribution() { curl https://sh.rustup.rs -sSf | sh -s -- -y --no-modify-path - export PATH=$RUSTUP_HOME/bin:$PATH - source $RUSTUP_HOME/env + # On Unix, rustup creates an env file. On Windows GitHub runners (MSYS bash), + # that file may not exist, so fall back to adding Cargo bin directly. + if [ -f "$CARGO_HOME/env" ]; then + # shellcheck disable=SC1090 + source "$CARGO_HOME/env" + elif [ -f "$RUSTUP_HOME/env" ]; then + # shellcheck disable=SC1090 + source "$RUSTUP_HOME/env" + else + export PATH="$CARGO_HOME/bin:$PATH" + fi # build and test rust @@ -123,18 +132,28 @@ test_source_distribution() { # Clone testing repositories into the expected location git clone https://github.com/apache/arrow-testing.git testing - git clone https://github.com/apache/parquet-testing.git parquet-testing + git clone https://github.com/apache/parquet-testing.git parquet python3 -m venv .venv - source .venv/bin/activate - python3 -m pip install -U pip - python3 -m pip install -U maturin - maturin develop + if [ -x ".venv/bin/python" ]; then + VENV_PYTHON=".venv/bin/python" + elif [ -x ".venv/Scripts/python.exe" ]; then + VENV_PYTHON=".venv/Scripts/python.exe" + elif [ -x ".venv/Scripts/python" ]; then + VENV_PYTHON=".venv/Scripts/python" + else + echo "Unable to find python executable in virtual environment" + exit 1 + fi + + "$VENV_PYTHON" -m pip install -U pip + "$VENV_PYTHON" -m pip install -U maturin + "$VENV_PYTHON" -m maturin develop #TODO: we should really run tests here as well #python3 -m pytest - if ( find -iname 'Cargo.toml' | xargs grep SNAPSHOT ); then + if ( find . -iname 'Cargo.toml' | xargs grep SNAPSHOT ); then echo "Cargo.toml version should not contain SNAPSHOT for releases" exit 1 fi diff --git a/docs/build.sh b/docs/build.sh index f73330323..e8409ee72 100755 --- a/docs/build.sh +++ b/docs/build.sh @@ -36,6 +36,21 @@ rm -rf build 2> /dev/null rm -rf temp 2> /dev/null mkdir temp cp -rf source/* temp/ + +# myst-nb executes each page as a notebook from the directory that page +# lives in, so the example data files must sit alongside every page that +# loads them by relative name (e.g. `ctx.read_csv("pokemon.csv")`). Symlink +# them into each directory that has such a page rather than copying the +# 20 MB parquet repeatedly. +for d in temp temp/user-guide temp/user-guide/common-operations; do + ln -sf "$script_dir/pokemon.csv" "$d/pokemon.csv" + ln -sf "$script_dir/yellow_tripdata_2021-01.parquet" "$d/yellow_tripdata_2021-01.parquet" +done + +# myst-nb runs `{code-cell}` blocks against a Jupyter kernel named "python3". +# Register the active environment's interpreter as that kernel (idempotent). +python -m ipykernel install --sys-prefix --name python3 --display-name "Python 3" + make SOURCEDIR=`pwd`/temp html cd "$original_dir" || exit 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/images/original_dark.svg b/docs/source/_static/images/original_dark.svg new file mode 100644 index 000000000..fbdf20ea7 --- /dev/null +++ b/docs/source/_static/images/original_dark.svg @@ -0,0 +1,31 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/docs/source/_static/theme_overrides.css b/docs/source/_static/theme_overrides.css index aaa40fba2..9eab5b03c 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); } @@ -91,3 +63,48 @@ Details: min(15vh, 110px) for the logo size, 8rem for search box etc*/ white-space: normal !important; } } + + +/* Hideable right-hand "On this page" sidebar. + * toc-toggle.js adds the button and toggles `pst-secondary-hidden` on ; + * hiding the sidebar lets the flex article container reclaim the width. */ + +body.pst-secondary-hidden .bd-sidebar-secondary { + display: none; +} + +/* Let the article use the freed space rather than just re-centering. */ +body.pst-secondary-hidden .bd-article-container { + max-width: none; +} + +/* Floating toggle button, pinned to the top-right under the navbar. */ +#pst-secondary-toggle { + position: fixed; + top: 4.5rem; + right: 0.75rem; + z-index: 1020; + display: flex; + align-items: center; + justify-content: center; + width: 2rem; + height: 2rem; + padding: 0; + border: 1px solid var(--pst-color-border, #ccc); + border-radius: 0.25rem; + background-color: var(--pst-color-surface, #fff); + color: var(--pst-color-text-base, #333); + cursor: pointer; +} + +#pst-secondary-toggle:hover { + color: rgb(var(--pst-color-link-hover)); +} + +/* The toggle is only meaningful where the sidebar is shown (wide screens); + * below the theme's lg breakpoint the sidebar is already collapsed away. */ +@media (max-width: 959.98px) { + #pst-secondary-toggle { + display: none; + } +} diff --git a/docs/source/_static/toc-toggle.js b/docs/source/_static/toc-toggle.js new file mode 100644 index 000000000..a5cc65e54 --- /dev/null +++ b/docs/source/_static/toc-toggle.js @@ -0,0 +1,66 @@ +/* + * 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. + */ + +/* Adds a button that hides the right-hand "On this page" sidebar so the + * article can use the full page width. The choice is remembered across + * pages via localStorage. */ +(function () { + "use strict"; + var KEY = "pst-secondary-hidden"; + + function apply(hidden, btn) { + document.body.classList.toggle("pst-secondary-hidden", hidden); + if (btn) { + btn.setAttribute("aria-pressed", String(hidden)); + btn.title = hidden ? "Show page contents" : "Hide page contents"; + } + } + + document.addEventListener("DOMContentLoaded", function () { + // Only offer the toggle on pages that actually have the sidebar. + if (!document.querySelector(".bd-sidebar-secondary")) { + return; + } + + var btn = document.createElement("button"); + btn.id = "pst-secondary-toggle"; + btn.type = "button"; + btn.setAttribute("aria-label", "Toggle page contents sidebar"); + btn.innerHTML = ''; + document.body.appendChild(btn); + + btn.addEventListener("click", function () { + var hidden = !document.body.classList.contains("pst-secondary-hidden"); + try { + localStorage.setItem(KEY, hidden ? "1" : "0"); + } catch (e) { + /* localStorage may be unavailable; toggle still works for this page. */ + } + apply(hidden, btn); + }); + + var stored = null; + try { + stored = localStorage.getItem(KEY); + } catch (e) { + /* ignore */ + } + apply(stored === "1", btn); + }); +})(); 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 01813b032..22bace809 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" @@ -48,16 +48,40 @@ extensions = [ "sphinx.ext.mathjax", "sphinx.ext.napoleon", - "myst_parser", - "IPython.sphinxext.ipython_directive", + # myst_nb is a superset of myst_parser: it provides the MyST markdown + # parser plus executable `{code-cell}` notebook directives. Do NOT also + # list "myst_parser" — myst_nb activates it internally and listing both + # raises an extension conflict. + "myst_nb", "autoapi.extension", ] +# NOTE: .rst stays alongside .md because sphinx-autoapi generates RST +# under autoapi/ and Sphinx needs the suffix to parse it. The human- +# authored docs are all MyST .md now. ".md" is routed through myst-nb so +# pages carrying jupytext/kernelspec front matter execute their +# `{code-cell}` blocks; pages without that front matter render as plain +# MyST markdown. The ".rst" entry is only for the autoapi build artifacts. source_suffix = { ".rst": "restructuredtext", - ".md": "markdown", + ".md": "myst-nb", } +# Execute notebook code cells at build time and fail the build if any cell +# raises — this replaces the old IPython sphinx directive, whose executed +# examples are now `{code-cell}` blocks. "force" re-executes every build so +# stale cached output can never ship. +nb_execution_mode = "force" +nb_execution_timeout = 120 +nb_execution_raise_on_error = True + +# Prefer the plain-text repr of a cell's last expression over its rich +# `_repr_html_`. A DataFrame's HTML repr is a self-contained widget (inline +# styles + an injected