如何使用Polars表达式替代map_groups自定义函数高效生成蛋白质序列的k-mer并添加分组特定索引？

我一直在研究Polars表达式库，想知道有没有更高效的方式实现蛋白质序列的k-mer生成——就是像下面这个kmerize函数做的事，把序列切成指定长度的子串：

def kmerize(sequence, ksize):
    kmers = [sequence[i : (i + ksize)] for i in range(len(sequence) - ksize + 1)]
    return kmers

目前我在Polars里的做法是先对sequence列分组，然后在map_groups里把序列重复对应长度的次数，炸开后添加行索引，再用索引去切片生成k-mer，最后过滤出长度正确的结果。

这里是一个最小可复现的例子：

from io import StringIO

import polars as pl

s = """sequence_name,sequence,length
sp|O43236|SEPT4_HUMAN Septin-4 OS=Homo sapiens OX=9606 GN=SEPTIN4 PE=1 SV=1,MDRSLGWQGNSVPEDRTEAGIKRFLEDTTDDGELSKFVKDFSGNASCHPPEAKTWASRPQVPEPRPQAPDLYDDDLEFRPPSRPQSSDNQQYFCAPAPLSPSARPRSPWGKLDPYDSSEDDKEYVGFATLPNQVHRKSVKKGFDFTLMVAGESGLGKSTLVNSLFLTDLYRDRKLLGAEERIMQTVEITKHAVDIEEKGVRLRLTIVDTPGFGDAVNNTECWKPVAEYIDQQFEQYFRDESGLNRKNIQDNRVHCCLYFISPFGHGLRPLDVEFMKALHQRVNIVPILAKADTLTPPEVDHKKRKIREEIEHFGIKIYQFPDCDSDEDEDFKLQDQALKESIPFAVIGSNTVVEARGRRVRGRLYPWGIVEVENPGHCDFVKLRTMLVRTHMQDLKDVTRETHYENYRAQCIQSMTRLVVKERNRNKLTRESGTDFPIPAVPPGTDPETEKLIREKDEELRRMQEMLHKIQKQMKENY,478
sp|O43521|B2L11_HUMAN Bcl-2-like protein 11 OS=Homo sapiens OX=9606 GN=BCL2L11 PE=1 SV=1,MAKQPSDVSSECDREGRQLQPAERPPQLRPGAPTSLQTEPQGNPEGNHGGEGDSCPHGSPQGPLAPPASPGPFATRSPLFIFMRRSSLLSRSSSGYFSFDTDRSPAPMSCDKSTQTPSPPCQAFNHYLSAMASMRQAEPADMRPEIWIAQELRRIGDEFNAYYARRVFLNNYQAAEDHPRMVILRLLRYIVRLVWRMH,198
sp|O60238|BNI3L_HUMAN BCL2/adenovirus E1B 19 kDa protein-interacting protein 3-like OS=Homo sapiens OX=9606 GN=BNIP3L PE=1 SV=1,MSSHLVEPPPPLHNNNNNCEENEQSLPPPAGLNSSWVELPMNSSNGNDNGNGKNGGLEHVPSSSSIHNGDMEKILLDAQHESGQSSSRGSSHCDSPSPQEDGQIMFDVEMHTSRDHSSQSEEEVVEGEKEVEALKKSADWVSDWSSRPENIPPKEFHFRHPKRSVSLSMRKSGAMKKGGIFSAEFLKVFIPSLFLSHVLALGLGIYIGKRLSTPSASTY,219
sp|O95197|RTN3_HUMAN Reticulon-3 OS=Homo sapiens OX=9606 GN=RTN3 PE=1 SV=2,MAEPSAATQSHSISSSSFGAEPSAPGGGGSPGACPALGTKSCSSSCADSFVSSSSSQPVSLFSTSQEGLSSLCSDEPSSEIMTSSFLSSSEIHNTGLTILHGEKSHVLGSQPILAKEGKDHLDLLDMKKMEKPQGTSNNVSDSSVSLAAGVHCDRPSIPASFPEHPAFLSKKIGQVEEQIDKETKNPNGVSSREAKTALDADDRFTLLTAQKPPTEYSKVEGIYTYSLSPSKVSGDDVIEKDSPESPFEVIIDKAAFDKEFKDSYKESTDDFGSWSVHTDKESSEDISETNDKLFPLRNKEAGRYPMSALLSRQFSHTNAALEEVSRCVNDMHNFTNEILTWDLVPQVKQQTDKSSDCITKTTGLDMSEYNSEIPVVNLKTSTHQKTPVCSIDGSTPITKSTGDWAEASLQQENAITGKPVPDSLNSTKEFSIKGVQGNMQKQDDTLAELPGSPPEKCDSLGSGVATVKVVLPDDHLKDEMDWQSSALGEITEADSSGESDDTVIEDITADTSFENNKIQAEKPVSIPSAVVKTGEREIKEIPSCEREEKTSKNFEELVSDSELHQDQPDILGRSPASEAACSKVPDTNVSLEDVSEVAPEKPITTENPKLPSTVSPNVFNETEFSLNVTTSAYLESLHGKNVKHIDDSSPEDLIAAFTETRDKGIVDSERNAFKAISEKMTDFKTTPPVEVLHENESGGSEIKDIGSKYSEQSKETNGSEPLGVFPTQGTPVASLDLEQEQLTIKALKELGERQVEKSTSAQRDAELPSEEVLKQTFTFAPESWPQRSYDILERNVKNGSDLGISQKPITIRETTRVDAVSSLSKTELVKKHVLARLLTDFSVHDLIFWRDVKKTGFVFGTTLIMLLSLAAFSVISVVSYLILALLSVTISFRIYKSVIQAVQKSEEGHPFKAYLDVDITLSSEAFHNYMNAAMVHINRALKLIIRLFLVEDLVDSLKLAVFMWLMTYVGAVFNGITLLILAELLIFSVPIVYEKYKTQIDHYVGIARDQTKSIVEKIQAKLPGIAKKKAE,1032
sp|P10415|BCL2_HUMAN Apoptosis regulator Bcl-2 OS=Homo sapiens OX=9606 GN=BCL2 PE=1 SV=2,MAHAGRTGYDNREIVMKYIHYKLSQRGYEWDAGDVGAAPPGAAPAPGIFSSQPGHTPHPAASRDPVARTSPLQTPAAPGAAAGPALSPVPPVVHLTLRQAGDDFSRRYRRDFAEMSSQLHLTPFTARGRFATVVEELFRDGVNWGRIVAFFEFGGVMCVESVNREMSPLVDNIALWMTEYLNRHLHTWIQDNGGWDAFVELYGPSMRPLFDFSWLSLKTLLSLALVGACITLGAYLGHK,239"""

ksize = 24

df = pl.scan_csv(StringIO(s))
df.group_by("sequence").map_groups(
    lambda group_df: group_df.with_columns(kmers=pl.col("sequence").repeat_by("length"))
    .explode("kmers")
    .with_row_index(),
    schema={"index": pl.UInt32, "sequence_name": pl.String, "sequence": pl.String, "length": pl.Int64, "kmers": pl.String},
).with_columns(pl.col("kmers").str.slice("index", ksize)).filter(
    pl.col("kmers").str.len_chars() == ksize
).rename(
    {"index": "start"}
).collect()

这段代码会生成这样的DataFrame：

不过我的数据集大概有25万条序列，每条长度在100到1000个字符之间，所以想找一种资源消耗更低的高效实现方式。有没有更好的办法呢？

谢谢大家，祝你们拥有美好的一天！

备注：内容来源于stack exchange，提问作者Olga Botvinnik