如何在R中用arrow包为Parquet数据集添加向量列并写入新文件？

问题重现 #

你遇到的核心问题是：用arrow::open_dataset读取Parquet数据集后，无法直接为其添加等长的向量列（单值列正常，多值向量报错），尝试DuckDB时也出现向量被逐行复用的错误，而全量加载内存的方法不适用于超大数据集。以下是你的测试代码及错误情况：

library(data.table)
library(arrow)
library(dplyr)
library(duckdb)

# 生成测试Parquet文件
dt = data.table(x = c(1, 2, 3), y = c("a", "b", "c"))
write_parquet(dt, "dt.parquet")

# 单值列可正常添加
ds <- open_dataset("dt.parquet")
ds |> mutate(z = 1) |> write_parquet("ds1.parquet")
read_parquet("ds1.parquet")
#       x      y     z
#   <num> <char> <num>
#1:     1      a     1
#2:     2      b     1
#3:     3      c     1

# 添加等长向量列报错
z = c(4,5,6)
ds |> mutate(z = z) |> write_parquet("ds2.parquet")
# 错误：In z = z, only values of size one are recycled
# Call collect() first to pull data into R.

# DuckDB尝试失败（向量被逐行复用为嵌套数组）
ds |> to_duckdb() |> mutate(z = z) |> to_arrow() |> write_parquet("ds2.parquet")
read_parquet("ds2.parquet")
# A tibble: 3 × 3
#      x y      z$``   $``   $``
#  <dbl> <chr> <dbl> <dbl> <dbl>
#1     1 a         4     5     6
#2     2 b         4     5     6
#3     3 c         4     5     6

解决方案 #

以下两种方法均无需全量加载数据集到内存，适配超大规模Parquet文件：

方法1：基于Arrow分块处理实现

利用Arrow的分块读取特性，为每个数据块匹配对应位置的向量切片，避免内存过载：

library(arrow)
library(dplyr)
library(purrr)

z <- c(4,5,6)
ds <- open_dataset("dt.parquet")

# 验证向量长度与数据集行数一致（无需加载全量数据）
total_rows <- ds |> summarize(n = n()) |> pull(n)
stopifnot(length(z) == total_rows)

# 定义分块处理函数：为每个分块截取对应位置的z向量
process_chunk <- function(chunk, chunk_idx, z_vec) {
  chunk_start <- (chunk_idx - 1) * chunk$num_rows + 1
  chunk_end <- chunk_idx * chunk$num_rows
  chunk_end <- min(chunk_end, length(z_vec)) # 处理最后一个分块的行数差异
  chunk_z <- z_vec[chunk_start:chunk_end]
  chunk |> mutate(z = chunk_z)
}

# 分块处理并写入新Parquet文件
ds |>
  chunked_array_reader() |>
  imap_dfr(process_chunk, z_vec = z) |>
  write_parquet("ds_arrow_chunked.parquet")

# 验证结果
read_parquet("ds_arrow_chunked.parquet")
# # A tibble: 3 × 3
#       x y         z
#   <dbl> <chr> <dbl>
# 1     1 a         4
# 2     2 b         5
# 3     3 c         6

方法2：正确使用DuckDB关联向量

将向量转为带行号的表，与原数据集按行号关联，避免向量被解析为嵌套数组：

library(arrow)
library(dplyr)
library(duckdb)

z <- c(4,5,6)
ds <- open_dataset("dt.parquet")

# 将z向量转为带行号的DuckDB表
z_tbl <- tibble(row_id = seq_along(z), z = z) |> to_duckdb()

# 原数据集添加行号后关联z表，再移除行号
ds |>
  to_duckdb() |>
  mutate(row_id = row_number()) |>
  left_join(z_tbl, by = "row_id") |>
  select(-row_id) |>
  to_arrow() |>
  write_parquet("ds_duckdb_fixed.parquet")

# 验证结果
read_parquet("ds_duckdb_fixed.parquet")
# # A tibble: 3 × 3
#       x y         z
#   <dbl> <chr> <dbl>
# 1     1 a         4
# 2     2 b         5
# 3     3 c         6

关键说明 #

• 方法1完全依托Arrow的分块机制，全程不会将数据集全量加载到内存，适合TB级别的超大型数据；
• 方法2利用DuckDB的SQL关联能力，逻辑更直观，同样无需全量加载原数据集；
• 两种方法都规避了collect()或as_arrow_table()带来的内存瓶颈。

内容的提问来源于stack exchange，提问作者julien.leroux5