andygrove opened a new pull request, #5144: URL: https://github.com/apache/datafusion-comet/pull/5144
## Which issue does this PR close? Backport of #5110 to `branch-1.0`. Part of #5109. ## Rationale for this change `SampleExec` is not supported on `branch-1.0` today, so any query using `DataFrame.sample`, SQL `TABLESAMPLE`, or `DataFrame.randomSplit` falls back to Spark and breaks up the surrounding native execution block. Spark has two sampling paths. Without replacement, it applies `BernoulliCellSampler` to each partition, seeded with `seed + partitionIndex`, drawing one `XORShiftRandom.nextDouble()` per input row and keeping the row when the value falls in `[lowerBound, upperBound)`. With replacement, it draws from a commons-math3 `PoissonDistribution` seeded by `Well19937c`. This covers the first path, which is the common case. `branch-1.0` already carries the bit-exact port of `XORShiftRandom` (used by `rand` / `randn` / `shuffle`), including the murmur3 `hashSeed` that `setSeed` applies, so reproducing Spark's exact draw sequence is a small addition rather than a new RNG port. The with-replacement path needs ports of `Well19937c` and `PoissonDistribution.sample()` and is left as follow-up work, tracked by #5109. ## What changes are included in this PR? Clean cherry-pick of 4b091ac with no conflicts and no adaptation needed: - New `Sample` protobuf message and `SampleExec` native operator (`native/core/src/execution/operators/sample.rs`), which filters each batch by the sampler's decisions and carries generator state across batch boundaries. - New `BernoulliCellSampler` in the `spark-expr` crate, a port of Spark's sampler built on the existing `XorShiftRandom`. - `CometSampleExec` operator serde, registered in `CometExecRule.nativeExecs`, gated by the new `spark.comet.exec.sample.enabled` config (default true). `getSupportLevel` returns `Unsupported` when `withReplacement` is set, so that case falls back to Spark. - The planner applies the partition index to the seed the same way the `rand` / `shuffle` expression builders do. - Documentation: a Sampling section in the operator compatibility guide covering the row-order caveat below, and a `SampleExec` row in the operator support reference. Because the sampler consumes one random value per row within a partition, it matches Spark's selection only when rows reach it in the same order. That holds when sampling above a scan, filter, or projection. Above an operator where Comet may emit rows in a different order than Spark, such as a join or an aggregate, the result is still a valid sample of the same expected size but not necessarily the same rows. This is inherent to per-row RNG consumption and already applies to `rand()`; it is documented in the compatibility guide. ## How are these changes tested? Rust unit tests cover that exactly one value is drawn per row, that complementary ranges partition the input, that an empty range selects nothing, and that splitting the input into multiple batches does not change which rows are selected. Tests in `CometExecSuite` compare results against Spark for `sample` with a fixed seed, for a multi-partition input (which is what catches a native side that ignores the partition index), and for `randomSplit` (non-zero lower bound), and assert that sampling with replacement and the disabled config both fall back to Spark. Verified on this branch that the native workspace passes `cargo check --all-targets --workspace` and that the JVM modules compile on both the default profile and `-Pspark-3.4`, which is the profile that exercises the shared `spark-3.x` shim. -- 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]
