[R] summarize after arrange fails

[r2evans]

Describe the bug, including details regarding any error messages, version, and platform.

If there's an arrange(.) in the lazy pipeline that is followed by some aggregation with summarize, the collection still looks for the sorting column:

library(arrow)
library(dplyr)
arrow_table(mtcars) |>
  summarize(across(mpg, list(Min = min, Max = max))) |>
  collect()
# # A tibble: 1 × 2
#   mpg_Min mpg_Max
#     <dbl>   <dbl>
# 1    10.4    33.9

arrow_table(mtcars) |>
  arrange(mpg) |>
  summarize(across(mpg, list(Min = min, Max = max))) |>
  collect()
# Error in compute.arrow_dplyr_query(x) : 
#   Invalid: Invalid sort key column: No match for FieldRef.Name(mpg) in mpg_Min: double
# mpg_Max: double
# ----
# mpg_Min:
#   [
#     [
#       10.4
#     ]
#   ]
# mpg_Max:
#   [
#     [
#       33.9
#     ]
#   ]

This example is somewhat contrived here, in that this summarization does not need ordered data. The underlying issue remains: why does it not sort the data at that point and then summarize? I'm not certain if this is a problem with lazy sorting or if it is too aggressive preserving the sort-field(s).

This behavior is in contrast to a selection removing the sorting column:

arrow_table(mtcars) |>
  arrange(mpg) |>
  select(-mpg) |>
  collect()
# # A tibble: 32 × 10
#      cyl  disp    hp  drat    wt  qsec    vs    am  gear  carb
#    <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl>
#  1     8  472    205  2.93  5.25  18.0     0     0     3     4
#  2     8  460    215  3     5.42  17.8     0     0     3     4
#  3     8  350    245  3.73  3.84  15.4     0     0     3     4
#  4     8  360    245  3.21  3.57  15.8     0     0     3     4
#  5     8  440    230  3.23  5.34  17.4     0     0     3     4
#  6     8  301    335  3.54  3.57  14.6     0     1     5     8
#  7     8  276.   180  3.07  3.78  18       0     0     3     3
#  8     8  304    150  3.15  3.44  17.3     0     0     3     2
#  9     8  318    150  2.76  3.52  16.9     0     0     3     2
# 10     8  351    264  4.22  3.17  14.5     0     1     5     4
# # ℹ 22 more rows
# # ℹ Use `print(n = ...)` to see more rows

> sessionInfo()

R version 4.4.2 (2024-10-31)
Platform: aarch64-apple-darwin20
Running under: macOS Sequoia 15.2

Matrix products: default
BLAS:   /Library/Frameworks/R.framework/Versions/4.4-arm64/Resources/lib/libRblas.0.dylib 
LAPACK: /Library/Frameworks/R.framework/Versions/4.4-arm64/Resources/lib/libRlapack.dylib;  LAPACK version 3.12.0

locale:
[1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8

time zone: America/New_York
tzcode source: internal

attached base packages:
[1] stats     graphics  grDevices utils     datasets  methods   base     

other attached packages:
[1] arrow_18.1.0.1 dplyr_1.1.4   

loaded via a namespace (and not attached):
 [1] assertthat_0.2.1 utf8_1.2.4       R6_2.5.1         bit_4.5.0.1      tidyselect_1.2.1 magrittr_2.0.3   glue_1.8.0       tibble_3.2.1     pkgconfig_2.0.3  bit64_4.5.2     
[11] generics_0.1.3   lifecycle_1.0.4  cli_3.6.3        fansi_1.0.6      vctrs_0.6.5      withr_3.0.2      compiler_4.4.2   purrr_1.0.2      pillar_1.9.0     rlang_1.1.4     

Component(s)

R

[r2evans]

This issue persists, and is not specific to just summarize, even a simple unrelated mutate will trigger the error (e.g., s/summarize/mutate/ above).

This seems a simple action, and the tie between the erroring columns and calculation is rather obscure. @thisisnic any thoughts on why this is happening, and lacking diving into apache C++ code, is there a way to fix or workaround this problem?

[thisisnic]

My initial thoughts are that it could be somewhere in the bits of query optimisation which are done in R before passing things to the C++ layer - CC @nealrichardson

I checked it wasn't something weird in the across() binding:

arrow_table(mtcars) |>
  arrange(mpg) |>
  summarize(min_mpg = min(mpg)) |>
  collect()

#> Error in `compute.arrow_dplyr_query()`:
#> ! Invalid: Invalid sort key column: No match for FieldRef.Name(mpg) in min_mpg: double
#> ----
#> min_mpg:
#>   [
#>     [
#>       10.4
#>     ]
#>   ]
#> Run `rlang::last_trace()` to see where the error occurred.

As a temporary workaround, you can, for some reason, call slice_head(n = nrow(data) - it appears to be forcing some ordering of operations, e.g.

library(arrow)
library(dplyr)

nrows <- nrow(mtcars)

arrow_table(mtcars) |>
  arrange(mpg) |>
  slice_head(n=nrows) |>
  summarize(across(mpg, list(Min = min, Max = max))) |>
  collect()
#> # A tibble: 1 × 2
#>   mpg_Min mpg_Max
#>     <dbl>   <dbl>
#> 1    10.4    33.9

Don't have time to delve deeply into it right now, but will point an LLM at it and see what it comes up with.

[thisisnic]

@r2evans - here's the summary it came up with. Looks reasonable to me - would you be interested in submitting a PR?

---

🤖 Generated with Claude Code

The Problem

When you call arrange() |> summarize(), the query stores the sort columns in arrange_vars. When the query is executed in ExecPlan$Build(), the aggregation happens first (producing only the aggregated columns like min_mpg), but then the code tries to apply sorting using the original arrange_vars which reference columns that no longer exist (like mpg).

The relevant code flow:

1. arrange(mpg) stores arrange_vars = {mpg} on the query
2. summarize() calls collapse() which nests the query - the inner query still has arrange_vars
3. In ExecPlan$Build() (query-engine.R), the aggregation is applied first (lines 100-129), then sorting is attempted (lines 164-178)
4. The sort fails because mpg no longer exists in the schema after aggregation

Why the slice_head() workaround works

slice_head() triggers its own collapse() with head set. When that inner query is built, both arrange_vars and head are set, so the sort is applied first, then the head is taken (that's the correct order for arrange() |> slice_head()).
The resulting outer query has empty arrange_vars, so when summarize() runs, there's nothing to sort.

The Fix

In do_arrow_summarize() (dplyr-summarize.R), clear arrange_vars before calling collapse():

# Clear arrange vars - sorting before aggregation is meaningless for Arrow's
# native aggregations, and the columns may not exist after aggregation.
# Users who want sorted output should arrange() after summarize(). (GH-45373)
.data$arrange_vars <- list()
.data$arrange_desc <- logical()

out <- collapse.arrow_dplyr_query(.data)

This is safe because:

• For ungrouped aggregation, sorting beforehand is semantically meaningless
• Arrow doesn't have native order-sensitive aggregations (e.g., first() falls back to R)
• If users want sorted output, they should arrange() after summarize()

[r2evans]

For the record: In the original example, we do lose the sort-column in the summarization, but if we mutate instead of summarize, we should still have the sort column so the sort is likely still appropriate (to the user, even if it may not affect the calculations).

arrow_table(mtcars) |>
  arrange(mpg) |>
  mutate(mpg_head = mean(mpg)) |>
  collect()
# Error in `compute.arrow_dplyr_query()`:
# ! Invalid: Invalid sort key column: No match for FieldRef.Name(mpg) in ..temp0: double
# ----
# ..temp0:
#   [
#     [
#       20.090625
#     ]
#   ]
# Run `rlang::last_trace()` to see where the error occurred.

Moving the arrange() to after mutate() works here, and it may work with the summarize() (if the code were slightly different ... unimportant atm).

Are you suggesting that it is best to silently drop the sort? If we cannot preserve $arrange_* during the collapse stage, some thoughts:

1. Drop the sort column(s) with a warning(). I think this is important, since the user needs to know that the assumption of ordered data is no longer true.
2. Is there a way to detect that this problem will occur and rearrange the steps such that arrange() occurs after the collapse? This presents the possibility that the column may not be available:
   • If it is not visible post-collapse (e.g., summarize() that drops the sort columns), then warn-and-drop;
   • If it is visible post-collapse, then rearrange the calc-steps so that the sorting is after the collapse-step(s).

BTW: I had used the slice_head() trick, but I ran into problems and had to remove it, whether for performance reasons or some other fragility. I can't find "why" I removed it from my code, so this is left dangling atm ...

Bottom line: while I understand that removing the sort in this case might be the hack to stop the error, I think it needs to be done with a warning instead of quietly. Thoughts?