nealrichardson commented on a change in pull request #10780: URL: https://github.com/apache/arrow/pull/10780#discussion_r675691031
########## File path: r/R/duckdb.R ########## @@ -0,0 +1,115 @@ +# 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. + +#' Create a (virtual) DuckDB table from an Arrow object +#' +#' This will do the necessary configuration to create a (virtual) table in DuckDB +#' that is backed by the Arrow object given. No data is copied or modified until +#' `collect()` or `compute()` are called or a query is run against the table. +#' +#' The result is a dbplyr-compatible object that can be used in d(b)plyr pipelines. +#' +#' Alternatively, one can pass the argument `.engine = "duckdb"` to `summarise()` +#' that starts with an Arrow object to use DuckDB to calculate the summarization +#' step. Internally, this calls `to_duckdb()` with all of the default argument +#' values. +#' +#' @param src the Arrow object (e.g. Dataset, Table) to use for the DuckDB table +#' @param con a DuckDB connection to use (default will create one and store it +#' in `options("arrow_duck_con")`) +#' @param table_name a name to use in DuckDB for this object. The default is a +#' unique string `"arrow_"` followed by numbers. +#' @param auto_disconnect should the table be automatically cleaned up when the +#' resulting object is removed (and garbage collected)? Default: `TRUE` +#' +#' @return A `tbl` of the new table in DuckDB +#' +#' @name to_duckdb +#' @export +#' @examplesIf arrow_with_dataset() && requireNamespace("duckdb", quietly = TRUE) && requireNamespace("dplyr", quietly = TRUE) +#' library(dplyr) +#' +#' ds <- InMemoryDataset$create(mtcars) +#' +#' ds %>% +#' filter(mpg < 30) %>% +#' to_duckdb() %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE)) +#' +#' # the same query can be simplified using .engine = "duckdb" +#' ds %>% +#' filter(mpg < 30) %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE), .engine = "duckdb") +#' +to_duckdb <- function( + src, Review comment: I might rename `src` to `.data` to be clear that we're working on the same dplyr inputs/objects. ########## File path: r/R/duckdb.R ########## @@ -0,0 +1,115 @@ +# 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. + +#' Create a (virtual) DuckDB table from an Arrow object +#' +#' This will do the necessary configuration to create a (virtual) table in DuckDB +#' that is backed by the Arrow object given. No data is copied or modified until +#' `collect()` or `compute()` are called or a query is run against the table. +#' +#' The result is a dbplyr-compatible object that can be used in d(b)plyr pipelines. +#' +#' Alternatively, one can pass the argument `.engine = "duckdb"` to `summarise()` +#' that starts with an Arrow object to use DuckDB to calculate the summarization +#' step. Internally, this calls `to_duckdb()` with all of the default argument +#' values. +#' +#' @param src the Arrow object (e.g. Dataset, Table) to use for the DuckDB table +#' @param con a DuckDB connection to use (default will create one and store it +#' in `options("arrow_duck_con")`) +#' @param table_name a name to use in DuckDB for this object. The default is a +#' unique string `"arrow_"` followed by numbers. +#' @param auto_disconnect should the table be automatically cleaned up when the +#' resulting object is removed (and garbage collected)? Default: `TRUE` +#' +#' @return A `tbl` of the new table in DuckDB +#' +#' @name to_duckdb +#' @export +#' @examplesIf arrow_with_dataset() && requireNamespace("duckdb", quietly = TRUE) && requireNamespace("dplyr", quietly = TRUE) +#' library(dplyr) +#' +#' ds <- InMemoryDataset$create(mtcars) +#' +#' ds %>% +#' filter(mpg < 30) %>% +#' to_duckdb() %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE)) +#' +#' # the same query can be simplified using .engine = "duckdb" +#' ds %>% +#' filter(mpg < 30) %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE), .engine = "duckdb") +#' +to_duckdb <- function( + src, + con = arrow_duck_connection(), + table_name = unique_arrow_tablename(), + auto_disconnect = TRUE) { + duckdb::duckdb_register_arrow(con, table_name, src) + + tbl <- tbl(con, table_name) + groups <- src$group + if (length(groups) > 0 && !is.null(groups)) { + tbl <- dplyr::group_by(tbl, !!rlang::sym(groups)) + } + + if (auto_disconnect) { + # this will add the correct connection disconnection when the tbl is gced. + # we should probably confirm that this use of src$disco is kosher. + tbl$src$disco <- duckdb_disconnector(con, table_name) + } + + tbl +} + +arrow_duck_connection <- function() { + con <- getOption("arrow_duck_con") + if (is.null(con) || !DBI::dbIsValid(con)) { + con <- DBI::dbConnect(duckdb::duckdb()) + # Use the same CPU count that the arrow library is set to + DBI::dbExecute(con, paste0("PRAGMA threads=", cpu_count())) + options(arrow_duck_con = con) + } + con +} + +# Adapted from dbplyr +unique_arrow_tablename <- function () { + i <- getOption("arrow_table_name", 0) + 1 + options(arrow_table_name = i) + sprintf("arrow_%03i", i) +} + +# Creates an environment that disconnects the database when it's GC'd +duckdb_disconnector <- function(con, tbl_name, quiet = FALSE) { Review comment: This arg doesn't seem to be used ```suggestion duckdb_disconnector <- function(con, tbl_name) { ``` ########## File path: r/R/duckdb.R ########## @@ -0,0 +1,115 @@ +# 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. + +#' Create a (virtual) DuckDB table from an Arrow object +#' +#' This will do the necessary configuration to create a (virtual) table in DuckDB +#' that is backed by the Arrow object given. No data is copied or modified until +#' `collect()` or `compute()` are called or a query is run against the table. +#' +#' The result is a dbplyr-compatible object that can be used in d(b)plyr pipelines. +#' +#' Alternatively, one can pass the argument `.engine = "duckdb"` to `summarise()` +#' that starts with an Arrow object to use DuckDB to calculate the summarization +#' step. Internally, this calls `to_duckdb()` with all of the default argument +#' values. +#' +#' @param src the Arrow object (e.g. Dataset, Table) to use for the DuckDB table +#' @param con a DuckDB connection to use (default will create one and store it +#' in `options("arrow_duck_con")`) +#' @param table_name a name to use in DuckDB for this object. The default is a +#' unique string `"arrow_"` followed by numbers. +#' @param auto_disconnect should the table be automatically cleaned up when the +#' resulting object is removed (and garbage collected)? Default: `TRUE` +#' +#' @return A `tbl` of the new table in DuckDB +#' +#' @name to_duckdb +#' @export +#' @examplesIf arrow_with_dataset() && requireNamespace("duckdb", quietly = TRUE) && requireNamespace("dplyr", quietly = TRUE) +#' library(dplyr) +#' +#' ds <- InMemoryDataset$create(mtcars) +#' +#' ds %>% +#' filter(mpg < 30) %>% +#' to_duckdb() %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE)) +#' +#' # the same query can be simplified using .engine = "duckdb" +#' ds %>% +#' filter(mpg < 30) %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE), .engine = "duckdb") +#' +to_duckdb <- function( + src, + con = arrow_duck_connection(), + table_name = unique_arrow_tablename(), + auto_disconnect = TRUE) { + duckdb::duckdb_register_arrow(con, table_name, src) Review comment: The first line in here should probably be `.data <- arrow_dplyr_query(.data)` to wrap/validate your input data. ########## File path: r/_pkgdown.yml ########## @@ -164,6 +164,9 @@ reference: - match_arrow - value_counts - list_compute_functions + - title: Connections to other systems + contents: + - to_duckdb Review comment: Scope creep: should we manufacture a man page for the r_to_py/py_to_r methods? ########## File path: r/R/duckdb.R ########## @@ -0,0 +1,115 @@ +# 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. + +#' Create a (virtual) DuckDB table from an Arrow object +#' +#' This will do the necessary configuration to create a (virtual) table in DuckDB +#' that is backed by the Arrow object given. No data is copied or modified until +#' `collect()` or `compute()` are called or a query is run against the table. +#' +#' The result is a dbplyr-compatible object that can be used in d(b)plyr pipelines. +#' +#' Alternatively, one can pass the argument `.engine = "duckdb"` to `summarise()` +#' that starts with an Arrow object to use DuckDB to calculate the summarization +#' step. Internally, this calls `to_duckdb()` with all of the default argument +#' values. +#' +#' @param src the Arrow object (e.g. Dataset, Table) to use for the DuckDB table +#' @param con a DuckDB connection to use (default will create one and store it +#' in `options("arrow_duck_con")`) +#' @param table_name a name to use in DuckDB for this object. The default is a +#' unique string `"arrow_"` followed by numbers. +#' @param auto_disconnect should the table be automatically cleaned up when the +#' resulting object is removed (and garbage collected)? Default: `TRUE` +#' +#' @return A `tbl` of the new table in DuckDB +#' +#' @name to_duckdb +#' @export +#' @examplesIf arrow_with_dataset() && requireNamespace("duckdb", quietly = TRUE) && requireNamespace("dplyr", quietly = TRUE) +#' library(dplyr) +#' +#' ds <- InMemoryDataset$create(mtcars) +#' +#' ds %>% +#' filter(mpg < 30) %>% +#' to_duckdb() %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE)) +#' +#' # the same query can be simplified using .engine = "duckdb" +#' ds %>% +#' filter(mpg < 30) %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE), .engine = "duckdb") +#' +to_duckdb <- function( Review comment: Can you fix the indentation of the args to be consistent with the other function definitions (which AFAIK follow the tidyverse style guide)? ########## File path: r/R/duckdb.R ########## @@ -0,0 +1,115 @@ +# 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. + +#' Create a (virtual) DuckDB table from an Arrow object +#' +#' This will do the necessary configuration to create a (virtual) table in DuckDB +#' that is backed by the Arrow object given. No data is copied or modified until +#' `collect()` or `compute()` are called or a query is run against the table. +#' +#' The result is a dbplyr-compatible object that can be used in d(b)plyr pipelines. +#' +#' Alternatively, one can pass the argument `.engine = "duckdb"` to `summarise()` +#' that starts with an Arrow object to use DuckDB to calculate the summarization +#' step. Internally, this calls `to_duckdb()` with all of the default argument +#' values. +#' +#' @param src the Arrow object (e.g. Dataset, Table) to use for the DuckDB table +#' @param con a DuckDB connection to use (default will create one and store it +#' in `options("arrow_duck_con")`) +#' @param table_name a name to use in DuckDB for this object. The default is a +#' unique string `"arrow_"` followed by numbers. +#' @param auto_disconnect should the table be automatically cleaned up when the +#' resulting object is removed (and garbage collected)? Default: `TRUE` +#' +#' @return A `tbl` of the new table in DuckDB +#' +#' @name to_duckdb +#' @export +#' @examplesIf arrow_with_dataset() && requireNamespace("duckdb", quietly = TRUE) && requireNamespace("dplyr", quietly = TRUE) +#' library(dplyr) +#' +#' ds <- InMemoryDataset$create(mtcars) +#' +#' ds %>% +#' filter(mpg < 30) %>% +#' to_duckdb() %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE)) +#' +#' # the same query can be simplified using .engine = "duckdb" +#' ds %>% +#' filter(mpg < 30) %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE), .engine = "duckdb") +#' +to_duckdb <- function( + src, + con = arrow_duck_connection(), + table_name = unique_arrow_tablename(), + auto_disconnect = TRUE) { + duckdb::duckdb_register_arrow(con, table_name, src) + + tbl <- tbl(con, table_name) + groups <- src$group + if (length(groups) > 0 && !is.null(groups)) { + tbl <- dplyr::group_by(tbl, !!rlang::sym(groups)) + } Review comment: If you do wrap in `arrow_dplyr_query`, or if you have received one, `$group` isn't where the group vars are stored (though partial matching will probably save you). If you don't wrap in that/have a Table/Dataset/RecordBatch, `$group` won't be what you want either. Reading our implementation for `dplyr::groups` I think this is what you want, not sure if the !! is needed. ```suggestion groups <- dplyr::groups(src) if (length(groups)) { tbl <- dplyr::group_by(tbl, !!groups) } ``` ########## File path: r/R/duckdb.R ########## @@ -0,0 +1,115 @@ +# 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. + +#' Create a (virtual) DuckDB table from an Arrow object +#' +#' This will do the necessary configuration to create a (virtual) table in DuckDB +#' that is backed by the Arrow object given. No data is copied or modified until +#' `collect()` or `compute()` are called or a query is run against the table. +#' +#' The result is a dbplyr-compatible object that can be used in d(b)plyr pipelines. +#' +#' Alternatively, one can pass the argument `.engine = "duckdb"` to `summarise()` +#' that starts with an Arrow object to use DuckDB to calculate the summarization +#' step. Internally, this calls `to_duckdb()` with all of the default argument +#' values. +#' +#' @param src the Arrow object (e.g. Dataset, Table) to use for the DuckDB table +#' @param con a DuckDB connection to use (default will create one and store it +#' in `options("arrow_duck_con")`) +#' @param table_name a name to use in DuckDB for this object. The default is a +#' unique string `"arrow_"` followed by numbers. +#' @param auto_disconnect should the table be automatically cleaned up when the +#' resulting object is removed (and garbage collected)? Default: `TRUE` +#' +#' @return A `tbl` of the new table in DuckDB +#' +#' @name to_duckdb +#' @export +#' @examplesIf arrow_with_dataset() && requireNamespace("duckdb", quietly = TRUE) && requireNamespace("dplyr", quietly = TRUE) +#' library(dplyr) +#' +#' ds <- InMemoryDataset$create(mtcars) +#' +#' ds %>% +#' filter(mpg < 30) %>% +#' to_duckdb() %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE)) +#' +#' # the same query can be simplified using .engine = "duckdb" +#' ds %>% +#' filter(mpg < 30) %>% +#' group_by(cyl) %>% +#' summarize(mean_mpg = mean(mpg, na.rm = TRUE), .engine = "duckdb") +#' +to_duckdb <- function( + src, + con = arrow_duck_connection(), + table_name = unique_arrow_tablename(), + auto_disconnect = TRUE) { + duckdb::duckdb_register_arrow(con, table_name, src) + + tbl <- tbl(con, table_name) + groups <- src$group + if (length(groups) > 0 && !is.null(groups)) { + tbl <- dplyr::group_by(tbl, !!rlang::sym(groups)) + } + + if (auto_disconnect) { + # this will add the correct connection disconnection when the tbl is gced. + # we should probably confirm that this use of src$disco is kosher. + tbl$src$disco <- duckdb_disconnector(con, table_name) + } + + tbl +} + +arrow_duck_connection <- function() { + con <- getOption("arrow_duck_con") + if (is.null(con) || !DBI::dbIsValid(con)) { + con <- DBI::dbConnect(duckdb::duckdb()) + # Use the same CPU count that the arrow library is set to + DBI::dbExecute(con, paste0("PRAGMA threads=", cpu_count())) + options(arrow_duck_con = con) + } + con +} + +# Adapted from dbplyr +unique_arrow_tablename <- function () { Review comment: nit ```suggestion unique_arrow_tablename <- function() { ``` ########## File path: r/tests/testthat/test-duckdb.R ########## @@ -0,0 +1,189 @@ +# 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. + +skip_if_not_installed("duckdb", minimum_version = "0.2.8") +skip_if_not_installed("dbplyr") +library(duckdb) +library(dplyr) + +con <- dbConnect(duckdb::duckdb()) +# we always want to test in parallel +dbExecute(con, "PRAGMA threads=2") Review comment: Why not set this in the `options(arrow_duck_connection)`? -- 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]
