Skip to content

Commit 5fa47b2

Browse files
committed
Merge remote-tracking branch 'apache/main' into andygrove/issue-6913-fix
Keep the scheduler probe from #6601 out of the way of a direct shuffle read: the probe's RDD hides the shuffle RDD that resolveInputObjects needs, so a write with a ShuffleScan leaf takes the shuffle RDD as it is and calls the probe's beforeInput at the top of CometIcebergWriteRDD.compute, ahead of the shuffle fetch. A child that is read over FFI still goes through the probe's RDD. Tests: the direct-read write test covers date, timestamp, timestamp_ntz, decimal, binary, boolean and float columns with NaN, -0.0, infinities and nulls, asserts that the coalesced case has a coalesced AQEShuffleReadExec, and adds a case with mostly empty reducer partitions. A new test checks that the probe runs for every task of a write that reads a shuffle directly.
2 parents fafef21 + 2f9552d commit 5fa47b2

618 files changed

Lines changed: 68865 additions & 14677 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.ai/skills/review-comet-docs/SKILL.md‎

Lines changed: 450 additions & 0 deletions
Large diffs are not rendered by default.

‎.ai/skills/review-comet-expression-pr/SKILL.md‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -225,9 +225,10 @@ and the requirement to record a performance audit. Hold the PR to both.
225225

226226
### Generated, review the source instead
227227

228-
The compatibility pages under `docs/source/user-guide/latest/compatibility/expressions/` and
229-
`docs/source/user-guide/latest/configs.md` are produced by `GenerateDocs` and regenerated by CI on
230-
merge. Never ask the contributor to hand-edit them. Check what feeds them:
228+
`GenerateDocs` fills in the expression compatibility tables and the config tables when the docs are
229+
built, and CI rebuilds and publishes the docs when a merge changes their inputs. The **Generated Content** section of
230+
`docs/source/contributor-guide/style_guide.md` says exactly which text is generated. Never ask the
231+
contributor to hand-edit that text. Check what feeds it:
231232

232233
- Expression marked `Incompatible(Some("..."))` but `getIncompatibleReasons()` is empty, so the
233234
compat page shows it as supported with no caveats.

‎.ai/skills/review-comet-memory-pr/SKILL.md‎

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -100,19 +100,27 @@ configuration, from the inside out:
100100

101101
- [ ] A new decorator forwards **every** `MemoryPool` method to its inner pool. A partial
102102
implementation makes `reserved()` disagree between levels.
103-
- [ ] **`fair_unified` checks the requesting consumer against its share, and the total against
104-
the pool.** It divides `pool_size` by the number of registered consumers and rejects if what
105-
the consumer holds plus the request exceeds that quotient, or if the pool's total reserved
106-
plus the request exceeds `pool_size`. With two consumers and an 8 GiB pool, once one holds
107-
3 GiB the other can still reserve up to 4 GiB. The pool keeps a running total for each
108-
consumer id, so the sibling reservations that `new_empty()`, `split()` and `take()` create
109-
count against one share. A PR that checks `reservation.size()` instead lets an operator with
110-
several reservations, such as a sort's streaming merge, take other consumers' shares. It
111-
also depends on when DataFusion updates the size, which is after it calls `try_grow` but
112-
before it calls `shrink`. If the PR changes either comparison it changes when every query
113-
spills, so it needs benchmark evidence, not reasoning.
114-
- [ ] **`num_consumers` counts every consumer in the task**, across every native plan, because the
115-
pool is task-shared.
103+
- [ ] **`fair_unified` holds the consumers that can spill to a share of the task's share.** The
104+
task's share is `spark.memory.offHeap.size` divided by the tasks on the executor running
105+
native plans, which `RunningTasks` counts while each task holds its task-shared pool, or
106+
`pool_size` if that is smaller. A consumer that can spill is
107+
refused if what it holds plus the request exceeds what the consumers that cannot spill leave
108+
of the task's share, divided by the registered consumers that can spill, or if the task's
109+
total plus the request exceeds the task's share. A consumer that cannot spill is refused only
110+
if the task's total plus the request exceeds `pool_size`, which binds before Spark's own
111+
limit only while the deprecated fraction is below `1 / N`. With eight running tasks, a 16 GiB
112+
pool and two consumers that can spill, each may hold 1 GiB, or 768 MiB while a merge holds
113+
512 MiB. The pool keeps a running total for each consumer id, so the sibling reservations
114+
that `new_empty()`, `split()` and `take()` create count against one share. A PR that checks
115+
`reservation.size()` instead lets an operator with several reservations, such as a sort's
116+
streaming merge, take other consumers' shares. It also depends on when DataFusion updates the
117+
size, which is after it calls `try_grow` but before it calls `shrink`. A PR that limits a
118+
consumer that cannot spill to a share, or sizes the shares from the whole pool, brings back
119+
#7088: one operator takes the task's share from Spark and a merge or hash join build fails
120+
the task. If the PR changes any of these comparisons it changes when every query spills, so
121+
it needs benchmark evidence, not reasoning.
122+
- [ ] **The consumers that divide the task's share are every consumer that can spill in the
123+
task**, across every native plan, because the pool is task-shared.
116124
- [ ] **Task-shared pool lifetime.** `acquire_task_shared_pool` keeps a process-wide
117125
`HashMap<task_attempt_id, Weak<TaskSharedMemoryPool>>`. The returned `Arc` is the only
118126
lifetime handle, so there is no explicit release to forget and a `createPlan` that fails
@@ -143,14 +151,16 @@ reclamation, that is a significant change and needs to say what it does to in-fl
143151
## 5. Budget and Configuration
144152

145153
`CometExecIterator.getMemoryConfig` computes the budget once per executor and passes it across JNI
146-
as `memory_limit`:
154+
as `memory_limit`, along with `spark.memory.offHeap.size` itself as `off_heap_size`, which
155+
`fair_unified` divides among the tasks running native plans:
147156

148157
```text
149158
memory_limit = spark.memory.offHeap.size * spark.comet.exec.memoryPool.fraction
159+
off_heap_size = spark.memory.offHeap.size
150160
```
151161

152162
- [ ] A new memory config goes through `getMemoryConfig` rather than being read separately in
153-
native code.
163+
native code, and a new value crossing JNI is added to `createPlan` on both sides.
154164
- [ ] Changing `memoryPool.fraction` semantics affects every deployment that tuned it as a haircut
155165
for the accounting gap.
156166
- [ ] On Kubernetes, `spark.memory.offHeap.size` is **part of** the pod limit, not headroom on top
@@ -204,8 +214,8 @@ will not catch. Check:
204214
- The **pool stack** diagram, if a decorator is added, removed, or reordered
205215
- The **pool type** table, if a pool type is added, removed, or changes how it is sized
206216
- The `memory_limit` formula, if the budget calculation changes
207-
- The **fair pool** description, which spells out the share and pool-total checks and their
208-
consequences
217+
- The **fair pool** description, which spells out the task share, the share and task-total
218+
checks, and their consequences
209219
- The **Crossing the FFI boundary** section, which names the operators that reserve for imported
210220
batches
211221
- The **Open problems** list. If the PR closes one of these gaps, the entry must be removed or

‎.ai/skills/review-comet-pr/SKILL.md‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
---
22
name: review-comet-pr
3-
description: Use when reviewing a DataFusion Comet pull request. Covers the workflow that applies to every PR and routes to the area-specific review skills for expressions, FFI, memory management, shuffle, and Iceberg writes. Provides guidance to a human reviewer rather than posting comments.
3+
description: Use when reviewing a DataFusion Comet pull request. Covers the workflow that applies to every PR and routes to the area-specific review skills for expressions, FFI, memory management, shuffle, Iceberg writes, and documentation. Provides guidance to a human reviewer rather than posting comments.
44
argument-hint: <pr-number>
55
---
66

@@ -43,6 +43,7 @@ change.
4343
| `native/core/src/execution/memory_pools/`, `CometTaskMemoryManager`, `CometShuffleMemoryAllocator`, `MemoryConsumer` / `try_grow` / `reserve` call sites, pool configs | `review-comet-memory-pr` |
4444
| `native/shuffle/`, `spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/`, `spark/src/main/java/org/apache/spark/shuffle/`, `CometShuffleExchangeExec` | `review-comet-shuffle-pr` |
4545
| `IcebergWriteStrategy`, `IcebergWriteExec`, `IcebergCommitExec`, `CometIcebergWriteExec`, `CometIcebergNativeWrite`, `iceberg_write.rs`, `iceberg_partition_path.rs`, the iceberg-rust pin | `review-comet-iceberg-write-pr` |
46+
| `docs/source/`, `README.md`, `AGENTS.md`, `.ai/skills/**/SKILL.md`, `.doc(...)` strings in config classes, doc comments | `review-comet-docs` |
4647

4748
For a new operator, read `docs/source/contributor-guide/adding_a_new_operator.md` alongside this
4849
skill. There is no dedicated operator review skill yet.
@@ -52,8 +53,8 @@ For a PR that deals with timestamps or the session timezone, read
5253
cuts across the serde, the native kernels, the scans and the JVM/native boundary, so no area skill
5354
owns it. "Timestamps and timezones" in step 5 says what to check.
5455

55-
If the PR falls outside all of these, for example build, CI, docs only, or release tooling, this
56-
skill on its own is the review.
56+
If the PR falls outside all of these, for example build, CI, or release tooling, this skill on its
57+
own is the review.
5758

5859
## 2. Gather PR Metadata
5960

@@ -95,7 +96,7 @@ Anything that reports `getSupportLevel()` must report it honestly:
9596

9697
- `Compatible()`, matches Spark exactly
9798
- `Incompatible(Some("reason"))`, differs in documented ways and requires
98-
`spark.comet.expr.allowIncompatible=true`
99+
`spark.comet.expression.<Name>.allowIncompatible=true`
99100
- `Unsupported(Some("reason"))`, cannot be implemented
100101

101102
A PR that quietly marks something `Compatible` while the diff shows a known divergence is the single
@@ -173,7 +174,7 @@ Report the comparison in the review even when nothing changed, so the reviewer k
173174

174175
Comet supports several Spark versions. The check on the Spark version belongs in Scala, either in a
175176
shim under `spark/src/main/spark-*/` or in one named predicate such as `isSpark42Plus` or
176-
`RegrSparkVersions.slopeFiltersVarByPairNulls`. It does not belong in an ad hoc comparison on the
177+
`SparkFixes.slopeFiltersVarByPairNulls`. It does not belong in an ad hoc comparison on the
177178
version string in shared code, or in native Rust. If the PR adds a shim for one 4.x version, check
178179
that the sibling 4.x source sets got it too.
179180

@@ -264,7 +265,8 @@ For a change on a default path, look for tests with the inputs that broke earlie
264265

265266
### CI
266267

267-
`gh pr checks <pr> --repo apache/datafusion-comet`, and again with `--failed` for detail.
268+
`gh pr checks <pr> --repo apache/datafusion-comet`. For only the failed checks, add
269+
`--json name,state,bucket,link --jq '.[] | select(.bucket == "fail")'`.
268270

269271
Summarize failures in the review. Do not compare against failures on `main`.
270272

@@ -299,10 +301,13 @@ same PR, because a follow-up that is not filed as an issue does not happen.
299301
Each area skill names the doc for its subsystem and the specific claims in it that go stale. When you
300302
load an area skill, read its doc and hold the diff against it.
301303

302-
**Do not ask contributors to hand-edit generated docs.** The compatibility guide pages under
303-
`docs/source/user-guide/latest/compatibility/expressions/` and
304-
`docs/source/user-guide/latest/configs.md` are produced by `GenerateDocs` and regenerated by CI on
305-
every merge to `main`. Review the source that feeds them instead.
304+
**Do not ask contributors to hand-edit generated docs.** `GenerateDocs` fills in the config tables
305+
and the expression compatibility tables when the docs are built, and CI rebuilds and publishes the
306+
docs when a merge to `main` changes their inputs. The
307+
**Generated Content** section of `docs/source/contributor-guide/style_guide.md` says exactly which
308+
text is generated. Review the source that feeds it instead. The prose around the generated tables,
309+
the `_category_template/` files, and the `index.md` pages are written by hand and are reviewed
310+
normally.
306311

307312
## Review Bar
308313

‎.asf.yaml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,9 @@ github:
116116
branch-1.1:
117117
required_pull_request_reviews:
118118
required_approving_review_count: 1
119+
branch-1.2:
120+
required_pull_request_reviews:
121+
required_approving_review_count: 1
119122
# GitHub's merge queue. Classic `protected_branches` above has no way to
120123
# express this, but .asf.yaml rulesets accept a raw Rulesets API payload, so
121124
# no INFRA ticket is needed. Same approach as apache/iceberg-python and
@@ -179,6 +182,13 @@ github:
179182
# Lets a committer queue a PR before its checks have finished, which is
180183
# most of the point of having a queue.
181184
allow_auto_merge: true
185+
# Limit open pull requests per user without write access.
186+
#
187+
# With coding agents, the number of open PRs now far exceeds our review capacity.
188+
# Setting a limit helps reviewers focus on fewer PRs at a time.
189+
creation_cap:
190+
enabled: true
191+
max_open_pull_requests: 3
182192
# publishes the content of the `asf-site` branch to
183193
# https://datafusion.apache.org/comet/
184194
publish:

‎.github/actions/rust-test/action.yaml‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,28 @@ runs:
2323
steps:
2424
# Note: cargo fmt check is now handled by the lint job that gates this workflow
2525

26+
- name: Set up embedded Python for native UDF tests
27+
shell: bash
28+
run: |
29+
apt-get update
30+
apt-get install -y --no-install-recommends python3 python3-dev python3-venv
31+
python3 -m venv /tmp/comet-rust-python
32+
/tmp/comet-rust-python/bin/pip install "pyarrow>=14" cloudpickle
33+
echo "PYO3_PYTHON=/tmp/comet-rust-python/bin/python" >> "$GITHUB_ENV"
34+
echo "PYTHONPATH=$(/tmp/comet-rust-python/bin/python -c 'import sysconfig; print(sysconfig.get_path("purelib"))')" >> "$GITHUB_ENV"
35+
2636
- name: Check Cargo clippy
2737
shell: bash
2838
run: |
2939
cd native
3040
cargo clippy --color=never --all-targets --workspace -- -D warnings
3141
42+
- name: Check native Python UDF feature
43+
shell: bash
44+
run: |
45+
cd native
46+
cargo clippy --color=never -p datafusion-comet --all-targets --features python-udf -- -D warnings
47+
3248
- name: Check compilation
3349
shell: bash
3450
run: |
@@ -83,6 +99,13 @@ runs:
8399
export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH}
84100
RUST_BACKTRACE=1 cargo nextest run
85101
102+
- name: Run native Python UDF tests
103+
shell: bash
104+
run: |
105+
cd native
106+
export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH}
107+
RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features python-udf python_udf
108+
86109
# The steps above lint and test the accounting allocator over the system allocator. Lint and
87110
# test it over jemalloc too, since that is a different allocator backend and nothing else in CI
88111
# builds it.

‎.github/dependabot.yml‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,42 @@ updates:
5353
update-types:
5454
- "minor"
5555
- "patch"
56+
- package-ecosystem: cargo
57+
directory: "/native/native-executors"
58+
schedule:
59+
interval: weekly
60+
target-branch: main
61+
labels: [dependencies]
62+
ignore:
63+
# major version bumps of datafusion*, arrow*, parquet are handled manually
64+
- dependency-name: "arrow*"
65+
update-types: ["version-update:semver-major"]
66+
- dependency-name: "parquet"
67+
update-types: ["version-update:semver-major"]
68+
- dependency-name: "datafusion*"
69+
update-types: ["version-update:semver-major"]
70+
groups:
71+
proto:
72+
applies-to: version-updates
73+
patterns:
74+
- "prost*"
75+
- "pbjson*"
76+
# Catch-all: group only minor/patch into a single PR,
77+
# excluding deps we want always separate (and excluding arrow/parquet which have their own group)
78+
all-other-cargo-deps:
79+
applies-to: version-updates
80+
patterns:
81+
- "*"
82+
exclude-patterns:
83+
- "arrow*"
84+
- "parquet"
85+
- "object_store"
86+
- "sqlparser"
87+
- "prost*"
88+
- "pbjson*"
89+
update-types:
90+
- "minor"
91+
- "patch"
5692
- package-ecosystem: "github-actions"
5793
directory: "/"
5894
schedule:

‎.github/labeler.yml‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -95,12 +95,9 @@ area:ffi:
9595
- changed-files:
9696
- any-glob-to-any-file:
9797
- 'native/jni-bridge/**'
98-
- 'native/core/src/execution/columnar_to_row.rs'
9998
- 'spark/src/main/**/vector/**'
10099
- 'common/src/main/**/vector/**'
101-
- 'spark/src/main/scala/org/apache/comet/NativeColumnarToRowConverter.scala'
102100
- 'spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala'
103-
- 'spark/src/main/scala/org/apache/spark/sql/comet/CometNativeColumnarToRowExec.scala'
104101
- 'spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala'
105102

106103
area:memory:

0 commit comments

Comments
 (0)