thisisnic opened a new pull request, #51385:
URL: https://github.com/apache/arrow/pull/51385
### Rationale for this change
`slice_sample()` was implemented with a UDF workaround because the C++
`random()` compute function could not be used inside an exec plan at the time
(GH-33181). That workaround has two problems:
1. On parquet datasets it usually returns 0 rows. The UDF takes
`is_null(<first column>)` as its input. When scanning parquet, the filter is
simplified against row-group statistics, and `is_null(col)` folds to a literal
`false` for row groups with no nulls. The UDF is then evaluated once per row
group with a length-1 input, so `runif(1) < prop` becomes a whole-row-group
keep/drop decision. With `n = 3` on `mtcars` that drops each row group about
86% of the time.
2. With `n`, the result is not a random sample. The old code oversampled
with a filter and then took `head(n)`, so the rows returned were almost always
the first survivors in scan order.
`random()` now works inside exec plans on Tables, Datasets and
RecordBatchReaders, in projections and in filters, so the UDF is no longer
needed.
### What changes are included in this PR?
`slice_sample()` is reimplemented on top of `Expression$create("random")`.
Design decisions, in rough order of importance:
- **`n` is implemented as sort-by-random then `head(n)`.** A fresh
`random()` column is added as a sort key, the query is sorted by it, and
`head(n)` is taken. This gives an exact, uniform sample without replacement.
The sort key is not a selected column, so `ensure_arrange_vars()` treats it as
a temporary column and the query engine projects it away after the `OrderBy`
node. Its name is made unique against both the selected columns and any
existing sort keys, since an existing sort key that is no longer selected also
becomes a temporary column.
- **Before sorting, the query is thinned with a random filter when the row
count is known.** Sorting materialises everything in the `OrderBy` node, which
is a problem when sampling a few thousand rows from hundreds of millions. When
`nrow()` is known, a filter `random() < oversample / nrows` is applied first
with `oversample = n + 10 * sqrt(n) + 100`. The number of surviving rows is
Binomial with standard deviation at most `sqrt(oversample) <= sqrt(n) + 10`, so
the margin of `10 * sqrt(n) + 100` is at least 10 standard deviations, making a
shortfall effectively impossible. Each row passes independently, so a uniform
sample of the survivors is a uniform sample of the whole. For `n = 10000` on
400M rows this sorts about 11k rows instead of 400M.
- **The row count is skipped when a source is a RecordBatchReader.**
Counting rows on a filtered query evaluates it, and a reader can only be read
once, so the later `collect()` would find it exhausted. In that case the
thinning step is skipped and the whole input is sorted. For filtered Tables and
Datasets the count is an extra pass over the data, but that was already true of
the old `n` path, and it is cheaper than sorting everything.
- **`n` no longer needs `nrow()` at all**, so it now works on joins,
aggregations and readers where it used to error. `n_to_prop()` is removed.
- **`prop` stays a streaming, approximate filter** (`random() < prop`),
matching the old behaviour apart from the random source. Converting `prop` to
an exact `n` when `nrow()` is known would match dplyr more closely, but would
turn `prop = 1` from a pass-through into a full sort and would lose streaming
for readers. The approximation is now documented.
- **The `seed` option of `RandomOptions` is deliberately not used.** The
kernel reseeds per batch, so a fixed seed would produce the identical sequence
in every batch. That means `set.seed()` no longer affects `slice_sample()` on
Arrow objects, which it did with the R UDF. This is documented.
- **Validation.** `n` must be a single non-negative number, `n` and `prop`
cannot both be given (dplyr errors too), and with neither given `n = 1` as in
dplyr.
- **`dim.arrow_dplyr_query()` fix.** For an unfiltered query over a
RecordBatchReader, `x$.data$num_rows` is NULL, so `dim()` collapsed to one
element and `nrow()` returned the column count. It now returns NA rows.
- **Cleanup.** The UDF, its `CanRunWithCapturedR()` requirement in tests,
the `runif` import and the docgen workaround for `_random_along` are removed. A
`query_has_reader()` helper replaces the duplicated "does this query read a
RecordBatchReader" check in `glimpse()`.
### Are these changes tested?
Yes. New tests cover: the partitioned parquet regression from the issue for
both `n` and `prop`; that samples are drawn from the whole table without
duplicates; `n` larger than the row count, `n = 0`, and no arguments; readers
with and without filters, including a wide reader that would expose the
`nrow()` bug; a user column named `..random`, alone and as a deselected sort
key; and input validation. Existing tests that only asserted "at most n rows"
now assert exactly `n`.
- [ ] TODO (@thisisnic): this code was written with AI assistance. Review
every line thoroughly before marking ready for review.
### Are there any user-facing changes?
Yes. `slice_sample()` now returns a uniformly random sample of exactly `n`
rows, including on parquet datasets where it previously returned 0 rows. `n`
now works on queries whose row count is not known up front. `set.seed()` no
longer affects results. Supplying both `n` and `prop` is now an error.
--
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]
