使用write_dataset()做Hive分区的factor列在open_dataset()中变为chr列的解决方法

我的数据包含一个存储为factor类型的Id列（值为字符与数字的组合），使用该列作为分区键，通过write_dataset()将数据以Hive风格分区存储为Arrow Parquet数据集。但通过open_dataset()读取数据时，Id列被转换为character类型。想请教两个问题：

1. 是否可以在继续用该列做分区的同时，让其自动转回factor类型？
2. 若无法实现，读取数据时手动转换为factor的最高效方法是什么？

演示代码：

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

idFactor <- factor(c(
  rep("A12-E123-I1234", 100),
  rep("B34-F456-J5678", 200),
  rep("C56-G789-K9101", 300),
  rep("D78-H101-L1121", 400)
))

factorData <- data.table(Id = idFactor, FactorId = idFactor)
testCols <- paste0("Column", seq(1, 2))
for (j in testCols) {
  set(factorData, j = j, value = runif(nrow(factorData), 0, 100000))
}

write_dataset(factorData, "factor_dataset", format = "parquet", partitioning = "Id", hive_style = TRUE)
ds <- open_dataset("factor_dataset")
blah <- ds %>% collect()

ds$schema # 架构显示Id为string，FactorId为预期的dictionary
str(blah) # str()显示Id为chr，FactorId为Factor

1. 能否自动转回factor类型？ #

不行。Hive风格的分区键本质是将分区值作为目录名的字符串存储，Arrow的open_dataset()读取分区列时，会默认将其解析为string（character）类型——分区目录里没有存储原factor的元信息（比如水平值、顺序），所以无法自动关联并转回factor类型。

2. 手动转换为factor的最高效方法 #

根据使用场景，推荐两种高效方案：

方案一：读取时指定自定义Schema（最优） #

提前定义包含factor对应类型的Schema，在open_dataset()中指定，读取时直接将Id列解析为factor，无需后续转换：

# 获取原数据的factor水平（也可手动指定）
id_levels <- levels(idFactor)
# 创建Arrow Schema：dictionary类型对应R的factor
custom_schema <- schema(
  Id = dictionary(string(), int32()),
  FactorId = dictionary(string(), int32()),
  Column1 = float64(),
  Column2 = float64()
)
# 用自定义Schema读取数据集
ds <- open_dataset("factor_dataset", schema = custom_schema)
blah <- ds %>% collect()
str(blah) # Id列会是Factor类型

方案二：Collect后用data.table高效修改 #

如果已经完成读取，用data.table的set()函数直接修改列类型是性能最优的选择（不复制整个数据集）：

ds <- open_dataset("factor_dataset")
blah <- ds %>% collect()
# 转为data.table并修改Id列
setDT(blah)
set(blah, j = "Id", value = factor(blah$Id, levels = id_levels))
str(blah)

如果习惯用dplyr，也可以用mutate(Id = factor(Id, levels = id_levels))，但大数据集下data.table的set()性能更优。

内容的提问来源于stack exchange，提问作者James