1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
|
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# The following S3 methods are registered on load if dplyr is present
group_by.arrow_dplyr_query <- function(.data,
...,
.add = FALSE,
add = .add,
.drop = dplyr::group_by_drop_default(.data)) {
.data <- as_adq(.data)
new_groups <- enquos(...)
# ... can contain expressions (i.e. can add (or rename?) columns) and so we
# need to identify those and add them on to the query with mutate. Specifically,
# we want to mark as new:
# * expressions (named or otherwise)
# * variables that have new names
# All others (i.e. simple references to variables) should not be (re)-added
# Identify any groups with names which aren't in names of .data
new_group_ind <- map_lgl(new_groups, ~ !(quo_name(.x) %in% names(.data)))
# Identify any groups which don't have names
named_group_ind <- map_lgl(names(new_groups), nzchar)
# Retain any new groups identified above
new_groups <- new_groups[new_group_ind | named_group_ind]
if (length(new_groups)) {
# now either use the name that was given in ... or if that is "" then use the expr
names(new_groups) <- imap_chr(new_groups, ~ ifelse(.y == "", quo_name(.x), .y))
# Add them to the data
.data <- dplyr::mutate(.data, !!!new_groups)
}
if (".add" %in% names(formals(dplyr::group_by))) {
# For compatibility with dplyr >= 1.0
gv <- dplyr::group_by_prepare(.data, ..., .add = .add)$group_names
} else {
gv <- dplyr::group_by_prepare(.data, ..., add = add)$group_names
}
.data$group_by_vars <- gv
.data$drop_empty_groups <- ifelse(length(gv), .drop, dplyr::group_by_drop_default(.data))
.data
}
group_by.Dataset <- group_by.ArrowTabular <- group_by.arrow_dplyr_query
groups.arrow_dplyr_query <- function(x) syms(dplyr::group_vars(x))
groups.Dataset <- groups.ArrowTabular <- function(x) NULL
group_vars.arrow_dplyr_query <- function(x) x$group_by_vars
group_vars.Dataset <- function(x) NULL
group_vars.RecordBatchReader <- function(x) NULL
group_vars.ArrowTabular <- function(x) {
x$r_metadata$attributes$.group_vars
}
# the logical literal in the two functions below controls the default value of
# the .drop argument to group_by()
group_by_drop_default.arrow_dplyr_query <-
function(.tbl) .tbl$drop_empty_groups %||% TRUE
group_by_drop_default.Dataset <- group_by_drop_default.ArrowTabular <-
function(.tbl) TRUE
ungroup.arrow_dplyr_query <- function(x, ...) {
x$group_by_vars <- character()
x$drop_empty_groups <- NULL
x
}
ungroup.Dataset <- force
ungroup.ArrowTabular <- function(x) {
x$r_metadata$attributes$.group_vars <- NULL
x
}
|
