andygrove commented on code in PR #6018: URL: https://github.com/apache/datafusion-comet/pull/6018#discussion_r4050391611
########## .ai/skills/review-comet-ffi-pr/SKILL.md: ########## @@ -0,0 +1,159 @@ +--- +name: review-comet-ffi-pr +description: Use when reviewing a DataFusion Comet pull request that crosses the JVM/native boundary, touching Arrow C Data or C Stream interface code, batch export and import, CometExecIterator, ScanExec, NativeUtil, CometVector subclasses, or jni_api. Load alongside review-comet-pr. +argument-hint: <pr-number> +--- + +<!-- +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. +--> + +FFI-specific review for Comet PR #$ARGUMENTS. + +**REQUIRED BACKGROUND:** Use `review-comet-pr` for PR metadata, existing comments, CI, the review +bar, and the output format. This skill only covers the JVM/native boundary. + +Bugs in this area do not produce wrong answers, they produce segfaults, leaks, and use-after-free +under load, often only on one platform or only when an operator buffers batches. Review it with +that in mind. + +## Read the Contributor Guide First + +| Doc | What you need from it | +| ---------------------------------------------------- | --------------------------------------------------------------------------- | +| `docs/source/contributor-guide/ffi.md` | Both data-flow directions, ownership rules, lifecycle, alignment workaround | +| `docs/source/contributor-guide/memory_management.md` | The "Crossing the FFI boundary" section: who is charged for a batch's bytes | + +Read `ffi.md` in full before the diff. The two directions have different ownership semantics and +reviewing one with the other's mental model is the most common way to miss a bug. + +## The Two Directions + +| Direction | Mechanism | Who owns the data | +| ----------------------------------- | ------------------------------------------ | ------------------------------------------------------------------ | +| JVM to native (`ScanExec`) | Arrow C **Stream**, one per partition | Native takes ownership by reference count when it imports a batch | +| Native to JVM (`CometExecIterator`) | Arrow C **Data**, one array pair per batch | Native allocates, JVM holds pointers and must `close()` to release | + +## 1. Ownership and Lifetime + +- [ ] **JVM to native: no defensive deep copies.** The C Stream transfers ownership by reference + count, so native can buffer imported batches in `SortExec` or the shuffle writer without + copying. A new `.clone()` of the data, a `copy_array`, or a "to be safe" deep copy on this + path is a real throughput cost. Ask what it is protecting against. +- [ ] **JVM to native: dropping the reader is the release.** When `ScanExec` drops its + `AlignedArrowStreamReader`, the stream's release callback fires synchronously back into the + JVM and closes the `ArrowReader` and its `VectorSchemaRoot`. Anything that extends the + reader's lifetime, stores it somewhere longer-lived, or drops it early changes when JVM + off-heap buffers are freed. +- [ ] **JVM to native: buffering pins JVM memory.** An operator that holds many imported batches + keeps the corresponding JVM-side off-heap buffers alive. A change that makes an operator + buffer more is also a memory change. +- [ ] **Native to JVM: every export needs a matching close.** Native allocates, the JVM wraps the + pointers in `ArrowBuf`s, and the bytes are freed only when the JVM calls `close()`. Trace the + new export to the `close()` that releases it, including on the exception path. +- [ ] **Release callbacks run on the error path too.** A batch exported and then abandoned because + the query failed still has to be released. Check the failure and cancellation paths, not just + the happy path. +- [ ] **No unwinding across `extern "C"`.** A Rust panic crossing the FFI boundary is undefined + behavior. New `#[no_mangle] extern "system"` entry points must not let a panic escape, and + must not `unwrap()` on anything an input can make fail. +- [ ] **Null and error checks on every pointer received from the other side.** + +## 2. The Stream Design + +The JVM exports each per-partition iterator **once** as an `ArrowArrayStream`, and native pulls +every batch through the stream's `get_next` callback. There is no per-batch JNI call and no +per-column FFI export on this path. + +A change that reintroduces per-batch or per-column export on the JVM-to-native path is a +performance regression even if it is correct. Flag it and ask why the stream could not carry it. + +The reader implementations in `CometNativeArrowSource.scala` are `RowArrowReader` for +`Iterator[InternalRow]`, `SparkColumnarArrowReader` for a non-Arrow `ColumnarBatch`, and +`ColumnarBatchArrowReader` for an Arrow-backed `ColumnarBatch`, which transfers `VectorSchemaRoot` +ownership. A new input shape needs a reader, not a special case elsewhere. + +## 3. Vector Types and Export Dispatch + +`NativeUtil.exportBatch()` matches on the concrete vector type. The `CometVector` hierarchy is +`CometDecodedVector` with Plain, Dictionary, List, Map, and Struct subclasses, plus +`CometSelectionVector` and `CometDelegateVector`. + +- [ ] A new `CometVector` subclass has a case in `exportBatch()` +- [ ] The case ordering is right. `CometSelectionVector` must be matched **before** the general + `CometVector` case, or the selection is silently dropped and the exported batch has the wrong + rows. +- [ ] Selection vectors are applied where `scan.rs` expects them, in `ScanExec::get_next()` Review Comment: Right, neither class exists anywhere in the tree. `exportBatch` has exactly two cases, `CometVector` and Spark's `ConstantColumnVector`, and `ScanExec` imports through `pull_next` with no selection hook at all, so both of those checks would have sent a reviewer chasing code that isn't there. Replaced the section in 67e734aad with the current model: the two export cases, the real `CometVector` hierarchy, and the point that a new subclass needs no new case as long as `getValueVector` returns an Arrow vector. The checks now cover who owns a materialized vector and the equal-value-count guard, which are the things that can actually go wrong on that path. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
