andygrove opened a new pull request, #5110: URL: https://github.com/apache/datafusion-comet/pull/5110
## Which issue does this PR close? Part of #5109. ## Rationale for this change `SampleExec` is not supported 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 PR covers the first path, which is the common case. Comet already has a 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? - 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? New 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. New 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. The result-comparison tests were verified to be non-vacuous by temporarily perturbing the native seed by one, which fails all three comparison tests. The suite passes on both the default profile and `-Pspark-3.4`. -- 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]
