rdplyrapache-arrow

Strategy form joining two very large arrow datasets without blowing up memory usage


I have two very large datasets in parquet files that I'm reading using R arrow::open(dataset). One file has over 20 million rows and the other, 15 million.

I need to join both datasets and save a new parquet file. So my first attempt was ver simple:

library(arrow)
library(dplyr)

ds1 <- open_dataset('part1.parquet/')
ds2 <- open_dataset('part2.parquet/')

all_data <- ds1 |>
  left_join(ds2, by = 'id')

write_dataset(all_data, 'all_data.parquet')

However, when attempting to write the data, my memory usage blows up (over 16Gb ram) and R crashes.

My second attempt was to break the work on a 'year' basis. And I was successful in saving a "joined" parquet file for each year, something on these lines:

library(arrow)
library(dplyr)

ds1 <- open_dataset('part1.parquet/')
ds2 <- open_dataset('part2.parquet/')

# r pseudo code here...
for (y in 2010:2020) {
  all_data <- ds1 |>
    filter(year == y) |>
    left_join(ds2, by = 'id')
  
  write_dataset(all_data, paste0('all_data_', y, '.parquet')
}

I then tried to merge those per year parquet files using open_dataset and dplyr::row_bind or `do.call(rbind...) but I was never able to get it working.

So the question is:

  1. Is there a better way to handle the merge such large datasets?

  2. If I have several parquet files that I open using arrow::open_dataset, how do I rbind them?

Sorry if I don't have a reprex. Not sure how to do one for such large dataset.

Also, this question seems related: Joining Arrow tables in R without overflowing memory or exceeding Acero's "bytes of key data" limit

> sessionInfo()
R version 4.2.1 (2022-06-23)
Platform: x86_64-pc-linux-gnu (64-bit)
Running under: Ubuntu 22.04 LTS

Matrix products: default
BLAS:   /usr/lib/x86_64-linux-gnu/blas/libblas.so.3.10.0
LAPACK: /usr/lib/x86_64-linux-gnu/lapack/liblapack.so.3.10.0

locale:
 [1] LC_CTYPE=en_US.UTF-8       LC_NUMERIC=C              
 [3] LC_TIME=en_US.UTF-8        LC_COLLATE=en_US.UTF-8    
 [5] LC_MONETARY=en_US.UTF-8    LC_MESSAGES=en_US.UTF-8   
 [7] LC_PAPER=en_US.UTF-8       LC_NAME=C                 
 [9] LC_ADDRESS=C               LC_TELEPHONE=C            
[11] LC_MEASUREMENT=en_US.UTF-8 LC_IDENTIFICATION=C       

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

other attached packages:
[1] dplyr_1.1.3    arrow_13.0.0.1

loaded via a namespace (and not attached):
 [1] fansi_1.0.5      utf8_1.2.4       assertthat_0.2.1 R6_2.5.1        
 [5] lifecycle_1.0.3  magrittr_2.0.3   pillar_1.9.0     rlang_1.1.1     
 [9] cli_3.6.1        vctrs_0.6.4      generics_0.1.3   bit64_4.0.5     
[13] glue_1.6.2       purrr_1.0.2      bit_4.0.5        compiler_4.2.1  
[17] pkgconfig_2.0.3  tidyselect_1.2.0 tibble_3.2.1    
> 

Solution

  • The key advantage of using arrow and parquet files is that you don't need to join them at the end - you can write a grouped dataframe and it will create a separate file for each group. In your loop, this would mean you can loop through all the years and write them all to the same dataset path whilst grouping by year first. Then the resulting datapath behaves like one joined dataframe:

    library(tidyverse)
    library(arrow)
    
    data_a <- tibble(
      id = rep(1:1000000, each = 10),
      year = rep(2011:2020, times = 1000000),
      a = sample(letters, 10000000, replace = TRUE)
    )
    
    data_b <- tibble(
      id = rep(1:1000000, each = 10),
      year = rep(2011:2020, times = 1000000),
      b = runif(10000000, 1, 100)
    )
    
    write_parquet(data_a, "data_a.parquet")
    write_parquet(data_b, "data_b.parquet")
    
    # Clean up
    rm(data_a)
    rm(data_b)
    
    # As if fresh start
    ds1 <- open_dataset("data_a.parquet")
    ds2 <- open_dataset("data_b.parquet")
    
    for (y in 2011:2020) {
      ds1 |>
        filter(year == y) |>
        left_join(ds2, by = c('id', 'year')) |> 
        group_by(year) |> 
        write_dataset("all_data")
    }
    
    # This creates a separate file for each yearly grouping, but treats as one dataset
    fs::dir_tree("all_data")
    #> all_data
    #> ├── year=2011
    #> │   └── part-0.parquet
    #> ├── year=2012
    #> │   └── part-0.parquet
    #> ├── year=2013
    #> │   └── part-0.parquet
    #> ├── year=2014
    #> │   └── part-0.parquet
    #> ├── year=2015
    #> │   └── part-0.parquet
    #> ├── year=2016
    #> │   └── part-0.parquet
    #> ├── year=2017
    #> │   └── part-0.parquet
    #> ├── year=2018
    #> │   └── part-0.parquet
    #> ├── year=2019
    #> │   └── part-0.parquet
    #> └── year=2020
    #>     └── part-0.parquet
    
    
    # As if fresh session - re-open dataset
    ds_joined <- open_dataset("all_data")
    
    # Doesn't need to bind rows. *should* be able to batch handle large datasets
    ds_joined |> 
      group_by(year) |> 
      summarise(min_b = min(b)) |> 
      collect()
    #> # A tibble: 10 × 2
    #>     year min_b
    #>    <int> <dbl>
    #>  1  2011  1.00
    #>  2  2012  1.00
    #>  3  2013  1.00
    #>  4  2014  1.00
    #>  5  2016  1.00
    #>  6  2017  1.00
    #>  7  2018  1.00
    #>  8  2015  1.00
    #>  9  2019  1.00
    #> 10  2020  1.00
    

    If you would be splitting your data up by year to do your analyses then this would be most suitable, as when you open and manipulate your dataset it will batch them in yearly groupings and only load the data needed for that year's calculation.

    Let me know if that doesn't solve your problem though. I've tested it here on a fairly column-light 10M row set of dataframes.