]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
1 | # Licensed to the Apache Software Foundation (ASF) under one |
2 | # or more contributor license agreements. See the NOTICE file | |
3 | # distributed with this work for additional information | |
4 | # regarding copyright ownership. The ASF licenses this file | |
5 | # to you under the Apache License, Version 2.0 (the | |
6 | # "License"); you may not use this file except in compliance | |
7 | # with the License. You may obtain a copy of the License at | |
8 | # | |
9 | # http://www.apache.org/licenses/LICENSE-2.0 | |
10 | # | |
11 | # Unless required by applicable law or agreed to in writing, | |
12 | # software distributed under the License is distributed on an | |
13 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
14 | # KIND, either express or implied. See the License for the | |
15 | # specific language governing permissions and limitations | |
16 | # under the License. | |
17 | ||
18 | ||
19 | # The following S3 methods are registered on load if dplyr is present | |
20 | ||
21 | arrange.arrow_dplyr_query <- function(.data, ..., .by_group = FALSE) { | |
22 | call <- match.call() | |
23 | exprs <- quos(...) | |
24 | if (.by_group) { | |
25 | # when the data is is grouped and .by_group is TRUE, order the result by | |
26 | # the grouping columns first | |
27 | exprs <- c(quos(!!!dplyr::groups(.data)), exprs) | |
28 | } | |
29 | if (length(exprs) == 0) { | |
30 | # Nothing to do | |
31 | return(.data) | |
32 | } | |
33 | .data <- as_adq(.data) | |
34 | # find and remove any dplyr::desc() and tidy-eval | |
35 | # the arrange expressions inside an Arrow data_mask | |
36 | sorts <- vector("list", length(exprs)) | |
37 | descs <- logical(0) | |
38 | mask <- arrow_mask(.data) | |
39 | for (i in seq_along(exprs)) { | |
40 | x <- find_and_remove_desc(exprs[[i]]) | |
41 | exprs[[i]] <- x[["quos"]] | |
42 | sorts[[i]] <- arrow_eval(exprs[[i]], mask) | |
43 | names(sorts)[i] <- format_expr(exprs[[i]]) | |
44 | if (inherits(sorts[[i]], "try-error")) { | |
45 | msg <- paste("Expression", names(sorts)[i], "not supported in Arrow") | |
46 | return(abandon_ship(call, .data, msg)) | |
47 | } | |
48 | descs[i] <- x[["desc"]] | |
49 | } | |
50 | .data$arrange_vars <- c(sorts, .data$arrange_vars) | |
51 | .data$arrange_desc <- c(descs, .data$arrange_desc) | |
52 | .data | |
53 | } | |
54 | arrange.Dataset <- arrange.ArrowTabular <- arrange.arrow_dplyr_query | |
55 | ||
56 | # Helper to handle desc() in arrange() | |
57 | # * Takes a quosure as input | |
58 | # * Returns a list with two elements: | |
59 | # 1. The quosure with any wrapping parentheses and desc() removed | |
60 | # 2. A logical value indicating whether desc() was found | |
61 | # * Performs some other validation | |
62 | find_and_remove_desc <- function(quosure) { | |
63 | expr <- quo_get_expr(quosure) | |
64 | descending <- FALSE | |
65 | if (length(all.vars(expr)) < 1L) { | |
66 | stop( | |
67 | "Expression in arrange() does not contain any field names: ", | |
68 | deparse(expr), | |
69 | call. = FALSE | |
70 | ) | |
71 | } | |
72 | # Use a while loop to remove any number of nested pairs of enclosing | |
73 | # parentheses and any number of nested desc() calls. In the case of multiple | |
74 | # nested desc() calls, each one toggles the sort order. | |
75 | while (identical(typeof(expr), "language") && is.call(expr)) { | |
76 | if (identical(expr[[1]], quote(`(`))) { | |
77 | # remove enclosing parentheses | |
78 | expr <- expr[[2]] | |
79 | } else if (identical(expr[[1]], quote(desc))) { | |
80 | # ensure desc() has only one argument (when an R expression is a function | |
81 | # call, length == 2 means it has exactly one argument) | |
82 | if (length(expr) > 2) { | |
83 | stop("desc() expects only one argument", call. = FALSE) | |
84 | } | |
85 | # remove desc() and toggle descending | |
86 | expr <- expr[[2]] | |
87 | descending <- !descending | |
88 | } else { | |
89 | break | |
90 | } | |
91 | } | |
92 | return( | |
93 | list( | |
94 | quos = quo_set_expr(quosure, expr), | |
95 | desc = descending | |
96 | ) | |
97 | ) | |
98 | } |