如何用itertools.product实现多列DataFrame与列表的笛卡尔积？

问题背景 #

给定如下DataFrame：

import pandas as pd
df = pd.DataFrame({'A': ['foo', 'bar', 'dex', 'tru'],
                   'B': ['abc', 'def', 'ghi', 'jkl']})

以及列表：

block_number = [1000, 2000]

需要生成包含原DataFrame所有列与列表元素的笛卡尔积组合，最终结果如下：

# 期望输出
pd.DataFrame({'A': ['foo', 'bar', 'dex', 'tru', 'foo', 'bar', 'dex', 'tru'],
              'B': ['abc', 'def', 'ghi', 'jkl', 'abc', 'def', 'ghi', 'jkl'],
              'block_number': [1000, 1000, 1000, 1000, 2000, 2000, 2000, 2000]})

现有基于itertools.product的方法仅支持单列DataFrame，需要适配多列甚至数千列的场景。

解决方案 #

方法1：利用merge实现笛卡尔积（推荐，适配任意列数） #

通过添加临时键列触发笛卡尔积合并，代码简洁且适合大量列的场景：

import pandas as pd

df = pd.DataFrame({'A': ['foo', 'bar', 'dex', 'tru'],
                   'B': ['abc', 'def', 'ghi', 'jkl']})
block_number = [1000, 2000]

# 将列表转为DataFrame
block_df = pd.DataFrame({'block_number': block_number})
# 添加临时键，实现笛卡尔积合并
df['tmp'] = 1
block_df['tmp'] = 1

# 合并后删除临时键，可选排序
result = df.merge(block_df, on='tmp').drop('tmp', axis=1)
result = result.sort_values('block_number').reset_index(drop=True)

print(result)

方法2：适配itertools.product处理多列 #

将DataFrame的行转为元组，再与列表元素生成组合，最后拆分还原列：

from itertools import product
import pandas as pd

df = pd.DataFrame({'A': ['foo', 'bar', 'dex', 'tru'],
                   'B': ['abc', 'def', 'ghi', 'jkl']})
block_number = [1000, 2000]

# 将每行转为元组
row_tuples = [tuple(row) for row in df.itertuples(index=False, name=None)]
# 生成所有行与block_number的组合
all_combinations = product(row_tuples, block_number)

# 转换为DataFrame并拆分元组列
result = pd.DataFrame(all_combinations, columns=['row_data', 'block_number'])
result[df.columns] = pd.DataFrame(result['row_data'].tolist(), index=result.index)
result = result.drop('row_data', axis=1).sort_values('block_number').reset_index(drop=True)

print(result)

方法3：重复原DataFrame并赋值（最直观） #

直接重复原DataFrame对应次数，再批量赋值block_number元素：

import pandas as pd

df = pd.DataFrame({'A': ['foo', 'bar', 'dex', 'tru'],
                   'B': ['abc', 'def', 'ghi', 'jkl']})
block_number = [1000, 2000]

# 重复原DataFrame len(block_number)次
result = pd.concat([df]*len(block_number), ignore_index=True)
# 生成block_number列：每个元素重复df行数的次数
result['block_number'] = [num for num in block_number for _ in range(len(df))]

print(result)

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