Spark2.3-SQL：基于指定字符串获取前后指定条数记录

嘿，针对你处理近200万条记录的需求，我整理了一个简洁高效的方案，完美适配你要根据搜索文本截取前后指定数量记录的场景，而且完全能应付大数据量的情况～

核心思路 #

先快速定位到匹配搜索文本的记录位置，再计算安全的截取范围（避免越界），最后直接切片提取结果——全程用Pandas的矢量化操作，效率拉满。

分步实现（适配你的场景） #

假设你的DataFrame是df，要搜索的列是label，目标文本是target，需要前后各取n条记录（比如你需求里的1000条，例子里的2条）：

1. 定位匹配项的索引

用布尔索引快速找到所有匹配的位置，这里先取第一个匹配项（如果有多个匹配，后面会讲怎么处理）：

# 找到所有匹配目标文本的索引
match_indices = df[df['label'] == target].index
# 取第一个匹配的索引（按需调整，比如取最后一个或全部）
match_idx = match_indices[0]

2. 计算安全的截取范围

为了避免匹配项在DataFrame的开头或结尾导致越界，用max和min约束起始/结束索引：

before_n = 1000  # 你需求里的前1000条，例子里改成2
after_n = 1000

# 如果是默认整数索引，用df.index[0]和df.index[-1]；如果是自定义索引，直接用对应边界
start_idx = max(df.index[0], match_idx - before_n)
end_idx = min(df.index[-1], match_idx + after_n)

3. 提取结果

用loc（自定义索引）或iloc（默认整数索引）直接切片，这一步是O(1)操作，完全不担心200万条数据的性能问题：

# 自定义索引用loc，默认整数索引用iloc
result = df.loc[start_idx:end_idx]

适配你的示例代码 #

把参数换成你给出的例子，运行后就能得到期望的输出：

import pandas as pd

# 构造你提供的示例DataFrame
data = {
    'index': [1,2,3,4,5,6,7,8,9,10],
    'X': [1,3,5,7,7,7,7,7,7,7],
    'label': ['A','B','C','D','E','F','G','H','I','J'],
    'date': ['2017-01-01','2017-01-02','2017-01-03','2017-01-04','2017-01-04','2017-01-04','2017-01-04','2017-01-04','2017-01-04','2017-01-04']
}
df = pd.DataFrame(data).set_index('index')  # 对齐你的示例索引

target = 'F'
before_n = 2
after_n = 2

match_indices = df[df['label'] == target].index
match_idx = match_indices[0]

start_idx = max(df.index[0], match_idx - before_n)
end_idx = min(df.index[-1], match_idx + after_n)

result = df.loc[start_idx:end_idx]
print(result)

输出结果：

X label        date
index                     
4      7     D  2017-01-04
5      7     E  2017-01-04
6      7     F  2017-01-04
7      7     G  2017-01-04
8      7     H  2017-01-04

处理多个匹配项的情况 #

如果搜索文本在DataFrame里有多个匹配，你可以合并所有匹配项的前后区间，避免重复记录：

all_ranges = []
for idx in match_indices:
    start = max(df.index[0], idx - before_n)
    end = min(df.index[-1], idx + after_n)
    all_ranges.append(range(start, end + 1))

# 合并所有区间的索引并去重
merged_indices = set()
for r in all_ranges:
    merged_indices.update(r)

# 按索引排序后提取结果
result = df.loc[sorted(merged_indices)]

性能说明 #

• 布尔索引找匹配项是Pandas的矢量化操作，比循环快几十倍，完全适合200万条数据的规模。
• 切片操作loc/iloc直接操作底层存储，不会复制整个DataFrame，内存占用极低。

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