如何使用R语言arrow包获取超大数据集的真实行索引

假设你有一个分割成多个片段、存储在磁盘上的超大文件，通过arrow包读取，示例代码如下：

X <- iris
library(arrow)
# bigX是存储在磁盘上的分块超大文件，不占用内存
bigX <- arrow_table(X)

用Sepal.Width == 3.8过滤后得到目标数据：

library(dplyr)
bigX %>% 
  filter(Sepal.Width == 3.8) %>% 
  collect()

得到的结果没有原数据集的行索引：

Sepal.Length Sepal.Width Petal.Length Petal.Width   Species
1          5.7         3.8          1.7         0.3    setosa
2          5.1         3.8          1.5         0.3    setosa
3          5.1         3.8          1.9         0.4    setosa
4          5.1         3.8          1.6         0.2    setosa
5          7.7         3.8          6.7         2.2 virginica
6          7.9         3.8          6.4         2.0 virginica

但需要像常规方法X[X$Sepal.Width==3.8,]那样，保留原完整数据集的真实行索引：

Sepal.Length Sepal.Width Petal.Length Petal.Width   Species
19           5.7         3.8          1.7         0.3    setosa
20           5.1         3.8          1.5         0.3    setosa
45           5.1         3.8          1.9         0.4    setosa
47           5.1         3.8          1.6         0.2    setosa
118          7.7         3.8          6.7         2.2 virginica
132          7.9         3.8          6.4         2.0 virginica

解决方法 #

Arrow的分块表默认不会保留原数据的行索引，要实现需求，需要在创建/存储分块文件时预先添加索引列，或者对已有的分块文件做预处理：

方法1：创建Arrow表时添加索引列 #

如果是从内存数据生成Arrow分块文件，先给原数据加上索引列，再转成arrow_table：

# 给原数据添加原行索引列，命名为original_index
X_with_index <- X %>% mutate(original_index = row_number())
# 生成带索引的Arrow分块表
bigX_with_index <- arrow_table(X_with_index)

之后过滤时保留该索引列，最后可以把它设为行名：

bigX_with_index %>% 
  filter(Sepal.Width == 3.8) %>% 
  collect() %>% 
  tibble::column_to_rownames(var = "original_index")

执行后就能得到带原行索引的结果：

Sepal.Length Sepal.Width Petal.Length Petal.Width   Species
19           5.7         3.8          1.7         0.3    setosa
20           5.1         3.8          1.5         0.3    setosa
45           5.1         3.8          1.9         0.4    setosa
47           5.1         3.8          1.6         0.2    setosa
118          7.7         3.8          6.7         2.2 virginica
132          7.9         3.8          6.4         2.0 virginica

方法2：处理已存储的Arrow分块文件 #

如果已经有存储在磁盘上的分块文件，没有索引列，需要先读取并添加索引后重新存储：

# 读取已有的Arrow分块文件
existing_bigX <- read_ipc_file("path/to/your/file")
# 添加索引列（注意：如果文件是分块的，row_number()会按分块顺序生成全局索引）
existing_bigX_with_index <- existing_bigX %>% mutate(original_index = row_number())
# 重新存储为带索引的分块文件
write_ipc_file(existing_bigX_with_index, "path/to/new/file_with_index")

之后就可以用方法1的过滤流程获取带原索引的结果。

注意事项 #

• 如果原数据是从外部文件（比如CSV、Parquet）导入生成Arrow分块表的，也要在导入时同步添加索引列，确保索引对应原数据集的行位置。
• 对于超大型文件，添加索引列的操作是在Arrow的计算引擎上执行的，不会占用过多内存，无需担心性能问题。

内容的提问来源于stack exchange，提问作者mr.T