[PR] GH-38638: [R] slice_sample returns 0 rows [arrow]

["thisisnic (via GitHub)" <gi...@apache.org>]

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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] GH-38638: [R] slice_sample returns 0 rows [arrow]

["github-actions[bot] (via GitHub)" <gi...@apache.org>]

github-actions[bot] commented on PR #51385:
URL: https://github.com/apache/arrow/pull/51385#issuecomment-5723367822

   :warning: GitHub issue #38638 **has been automatically assigned in GitHub** to PR creator.


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org