在SageMaker Studio中用PyAthena抽取随机1%样本并添加样本权重的问题

问题原因 #

Athena不支持在LIMIT子句中使用动态表达式或子查询，这是触发InvalidRequestException的直接原因。

方法一：使用TABLESAMPLE快速近似抽样 #

适合大数据集，引擎层面直接实现近似1%的随机抽样，执行效率高，同时可直接添加样本权重字段：

SELECT 
    *,
    100 AS sample_weight  -- 抽样比例1%，权重为1/0.01=100
FROM myDataMart
TABLESAMPLE BERNOULLI(1)  -- 按1%概率随机选取每条记录

对应PyAthena代码 #

from pyathena import connect

s3_query_results = 'my s3 Location'
region = 'eu-west-2'
workgroup='primary'

# 构造抽样SQL
sql_query = '''
SELECT 
    *,
    100 AS sample_weight
FROM myDataMart
TABLESAMPLE BERNOULLI(1)
'''

# 执行查询
cursor = connect(s3_staging_dir=s3_query_results, region_name=region, work_group=workgroup).cursor()
cursor.execute(sql_query)

方法二：精确1%随机抽样（分两步执行） #

如果需要严格匹配1%的样本量，需先获取总记录数计算精确LIMIT值，再执行随机抽样：

步骤1：查询数据集总记录数 #

# 初始化连接
cursor = connect(s3_staging_dir=s3_query_results, region_name=region, work_group=workgroup).cursor()

# 获取总记录数
count_query = "SELECT COUNT(*) FROM myDataMart"
cursor.execute(count_query)
total_count = cursor.fetchone()[0]
sample_size = int(total_count * 0.01)

步骤2：执行精确随机抽样并添加权重 #

# 构造精确抽样SQL
sql_query = f'''
SELECT 
    *,
    100 AS sample_weight
FROM myDataMart
ORDER BY RAND()
LIMIT {sample_size}
'''

# 执行抽样查询
cursor.execute(sql_query)

说明 #

• sample_weight的计算逻辑：抽取1%样本意味着每条样本对应原数据集中100条记录，因此权重设为100。
• 方法一为近似抽样，速度快；方法二是精确抽样，需两次查询，适合对样本量精度要求高的场景。

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