按 chunk_size 写入多个 parquet 文件答案

【问题标题】：write multiple parquet files by chunk_size按 chunk_size 写入多个 parquet 文件
【发布时间】：2021-12-24 04:26:40
【问题描述】：

我在arrow::write_parquet() 中看到了chunk_size 参数，但它的行为似乎不像预期的那样。我希望下面的代码生成 3 个单独的镶木地板文件，但只创建一个，并且 nrow > chunk_size。

library(arrow)
# .parquet dir and file path
td <- tempdir()
tf <- tempfile("", td, ".parquet")
on.exit(unlink(tf))

# dataframe with 3e6 rows
n  <- 3e6 
df <- data.frame(x = rnorm(n))

# write with chunk_size 1e6, and view directory
write_parquet(df, tf, chunk_size = 1e6)
list.files(td)

返回一个文件而不是 3 个：

[1] "25ff74854ba6.parquet"

# read parquet and show all rows are there
nrow(read_parquet(tf))

[1] 3000000

我们无法将多个文件名参数传递给write_parquet()，而且我不想分区，所以write_dataset() 似乎也不适用。

【问题讨论】：

标签： r parquet apache-arrow

【解决方案1】：

chunk_size 参数是指一次写入磁盘的数据量，而不是生成的文件数。 write_parquet() 函数旨在写入单个文件，而正如您所说，write_dataset() 允许分区文件写入。我认为目前不支持在任何其他基础上拆分文件，尽管在未来的版本中可能会这样做。如果您有特定原因需要 3 个单独的文件，我建议您先将数据分成多个数据集，然后通过 write_parquet() 编写每个数据集。

（另外，我是 R 包的开发人员之一，从文档中可以看出这一点并不完全清楚，所以我要开一张票来更新这些 - 感谢您标记这一点）

【讨论】：

谢谢@thisisnic。这很有意义，我同意这是对文档的快速修复。我认为如果数据集足够大，能够将多个较小的 parquet 文件分块以避免创建大量压缩文件将是一个有用的功能。传送文件要容易一些，它们可以用open_dataset() 和collect() 打开。我将在下面发布一个答案，说明我是如何解决这个问题的。

【解决方案2】：

write_parquet() 的参数不足，如 max_row，默认为合理的数字（如 1e6），我们可以这样做：

library(arrow)
library(uuid)
library(glue)
library(dplyr)

write_parquet_multi <- function(df, dir_out, max_row = 1e6){
  
  # Only one parquet file is needed
  if(nrow(df) <= max_row){
    cat("Saving", formatC(nrow(df), big.mark = ","), 
        "rows to 1 parquet file...")
    write_parquet(
      df, 
      glue("{dir_out}/{UUIDgenerate(use.time = FALSE)}.parquet"))
    cat("done.\n")
  }
  
  # Multiple parquet files are needed
  if(nrow(df) > max_row){
    count = ceiling(nrow(df)/max_row)
    start = seq(1, count*max_row, max_row)
    end   = c(seq(max_row, nrow(df), max_row), nrow(df))
    uuids = UUIDgenerate(n = count, use.time = FALSE)
    
    cat("Saving", formatC(nrow(df), big.mark = ","), 
        "rows to", count, "parquet files...")
    for(j in 1:count){
      write_parquet(
        dplyr::slice(df, start[j]:end[j]), 
        glue("{dir_out}/{uuids[j]}.parquet"))
    }
    cat("done.\n")
  }
  
}


# .parquet dir and file path
td <- tempdir()
tf <- tempfile("", td, ".parquet")
on.exit(unlink(tf))

# dataframe with 3e6 rows
n  <- 3e6 
df <- data.frame(x = rnorm(n))

# write parquet multi
write_parquet_multi(df, td)

list.files(td)

这会返回：

[1] "7a1292f0-cf1e-4cae-b3c1-fe29dc4a1949.parquet"    
[2] "a61ac509-34bb-4aac-97fd-07f9f6b374f3.parquet"    
[3] "eb5a3f95-77bf-4606-bf36-c8de4843f44a.parquet"

【讨论】：