DuckDB表抽样函数问题：重复执行报错且结果不一致

问题背景 #

需要实现一个函数对DuckDB表执行两种抽样操作：

1. 按比例抽取表中的行
2. 按比例抽取指定列的唯一条目，返回包含这些条目的所有行

当前函数单次运行正常，但重复执行时出现两个问题：

• 偶尔触发“表已存在”错误（即使设置了overwrite=TRUE）
• 已设置随机种子，但每次输出结果仍不一致

原示例数据集 #

# 创建数据框
id <- sample(c(1:100, NA), 150, replace = TRUE)
value <- sample(c(1:100, NA), 150, replace = TRUE)
sampletst_df <- data.frame(id = id, value = value)

# 转换为DuckDB表
tstcon <- dbConnect(duckdb::duckdb())
dbWriteTable(tstcon, "sample", sampletst_df, overwrite = TRUE)
samptst_dk <- tbl(tstcon,  "sample")

原函数实现 #

library(Duckdb)
library(DBI)

sampfrac.dk <- function(table_name, frac, var_name) {
  if (frac < 0 || frac > 1) {
    stop("The frac input must be between 0 and 1.")
  }
  df <- (table_name %>% ungroup() %>% summarise(n = n()) %>% collect())
  seed <- as.numeric(df[1, 1])
  set.seed(seed)
  con <- dbConnect(duckdb::duckdb())
  if (missing(var_name)) {
    # 获取行数
    row_cnt_df <- table_name %>% summarise(n = n()) %>% collect()
    row_cnt <- as.numeric(row_cnt_df[1, 1])
    row_num <- round(frac * row_cnt)
    # 创建抽样数据框
    df <- data.frame(row_num = sample(1:row_cnt, row_num, replace = FALSE))
    # 将抽样数据框写入DuckDB
    dbWriteTable(con, "frac_data", df, overwrite = TRUE)
    # 执行内连接
    samp_dk <- tbl(con, "frac_data")
    joined_data <- table_name %>%
      mutate(row_num = row_number()) %>%
      inner_join(samp_dk, by = "row_num", copy = TRUE) %>%
      arrange(row_num)
    dbDisconnect(con, shutdown = TRUE)
    return(joined_data)
  } else {
    # 创建唯一变量集合
    var <- deparse(substitute(var_name))
    dis_tbl <- table_name %>%
      select(all_of(var)) %>%
      distinct()
    # 获取行数
    row_cnt_df <- table_name %>% summarise(n = n()) %>% collect()
    row_cnt <- as.numeric(row_cnt_df[1, 1])
    row_num <- round(frac * row_cnt)
    # 保存为数据框
    df <- data.frame(row_num = sample(1:row_cnt, row_num, replace = FALSE))
    # 写入DuckDB表
    dbWriteTable(con, "frac_data", df, overwrite = TRUE)
    dsamp_dk <- tbl(con, "frac_data")
    # 创建抽样数据集
    samp_dk <- dis_tbl %>%
      mutate(row_num = row_number()) %>%
      inner_join(dsamp_dk, by = "row_num", copy = TRUE) %>%
      select(all_of(var))
    # 关联抽样数据与主数据
    joined_data <- table_name %>%
      inner_join(samp_dk, copy = TRUE)  %>%
      arrange(row_num)
    dbDisconnect(con, shutdown = TRUE)
    return(joined_data)
  }
}

问题分析 #

1. 连接冲突：函数内每次新建DuckDB连接，与原表所在的外部连接属于不同数据库实例，overwrite=TRUE仅对当前连接有效，跨实例操作易引发表存在性错误。
2. 种子无效：原代码用表行数作为种子，但row_number()在DuckDB端执行，与R端随机数生成不同步，导致结果无法复现；同时列抽样错误使用主表行数计算抽样量，而非去重后的列条目数。
3. 冗余操作：手动创建抽样索引表的流程冗余，且增加了出错概率。

修复后的函数 #

library(duckdb)
library(DBI)
library(dplyr)

sampfrac.dk <- function(table_tbl, frac, var_name) {
  # 校验抽样比例范围
  if (frac < 0 || frac > 1) {
    stop("抽样比例必须在0到1之间")
  }
  
  # 复用原表的连接，避免跨实例冲突
  con <- table_tbl$src$con
  
  # 设置同步的随机种子，确保结果可复现
  set.seed(123) # 可改为自定义参数让用户传入
  duckdb_seed <- as.integer(runif(1, 0, 2^31 - 1))
  dbExecute(con, paste0("SET seed = ", duckdb_seed))
  
  if (missing(var_name)) {
    # 行抽样：直接使用dplyr内置的sample_frac，底层调用DuckDB原生逻辑
    sampled <- table_tbl %>%
      sample_frac(frac, replace = FALSE) %>%
      arrange(row_number())
  } else {
    # 指定列条目抽样：先抽取唯一列值，再关联原表
    var <- ensym(var_name)
    
    # 计算去重后的列条目数，确定抽样数量
    distinct_cnt <- table_tbl %>%
      select(!!var) %>%
      distinct() %>%
      summarise(n = n()) %>%
      pull()
    
    sample_n <- round(frac * distinct_cnt)
    
    # 抽取目标列的唯一值
    sampled_vars <- table_tbl %>%
      select(!!var) %>%
      distinct() %>%
      sample_n(sample_n, replace = FALSE)
    
    # 关联原表，返回所有包含抽样条目的行
    sampled <- table_tbl %>%
      inner_join(sampled_vars, by = as.character(var)) %>%
      arrange(!!var)
  }
  
  return(sampled)
}

关键优化点 #

• 复用连接：直接从传入的tbl对象获取已有连接，避免跨实例操作的冲突，提升运行效率。
• 同步种子：同时设置R端和DuckDB端的随机种子，确保抽样结果完全可复现。
• 修正抽样逻辑：列抽样时基于去重后的列条目数计算抽样量，符合“按比例抽取列条目”的需求。
• 简化流程：使用dplyr内置的抽样函数，无需手动创建中间表，彻底解决表存在性问题。

测试代码 #

# 测试行抽样
row_sample <- sampfrac.dk(samptst_dk, 0.9)
print(row_sample)

# 测试指定列抽样（按id列）
col_sample <- sampfrac.dk(samptst_dk, 0.5, id)
print(col_sample)

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