能否通过移入父目录加.zgroup文件合并按染色体的Zarr存储？

我通过SLURM任务数组启动Python程序，利用scikit-allel从多样本VCF文件生成了按染色体划分的Zarr存储，以此并行化作业来缩短耗时、降低内存占用。现在能否通过重命名这些Zarr存储、将其移入同一父目录并添加.zgroup文件，将14个Zarr存储合并为单个Zarr存储？以下是原实现代码：

#!/usr/bin/env python3

import sys; print(sys.version)
import os
import glob
import subprocess
import numpy as np; print('numpy', np.__version__)
import pandas as pd; print('pandas',pd.__version__)
import allel; print('allel', allel.__version__)
import zarr; print('zarr', zarr.__version__)

INFN = sys.argv[1]
if not INFN:
    print('Must provide input .vcf.gz as first argument')
    sys.exit(2)

FIELDS = [
    'samples',
    'variants/CHROM',
    'variants/POS',
    'variants/REF',
     'variants/ALT',
    'variants/QUAL',
    'variants/TYPE',
    'variants/is_snp',
    'variants/numalt',
    'variants/AF',
    'variants/DP',
     'variants/ANN',
    'calldata/DP',
    'calldata/GT',
         ]
EXCLUDE_FIELDS = None

TABIX_EXEC = 'tabix'

print("Using tabix executable '{}' {} '{}'\n{}".format(TABIX_EXEC, "-",
        subprocess.check_output(['which', 'tabix']).decode('utf-8').rstrip(),
        subprocess.check_output([TABIX_EXEC, '--version']).decode('utf-8')))

task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))

chroms = subprocess.check_output([TABIX_EXEC,'-l',INFN],
                                 universal_newlines=True).strip().split('\n')

ch = chroms[task_id]

OUTFN = f"{INFN}.{ch}.zarr"

transformers = None
if 'ANN' in FIELDS:
    transformers=allel.ANNTransformer()

def vcf_to_zarr_func(ch):
    allel.vcf_to_zarr(INFN, OUTFN,
                      region=ch,
                      group=ch,
                      log=sys.stderr,
                      fields=FIELDS,
                      exclude_fields=EXCLUDE_FIELDS,
                      tabix=TABIX_EXEC,
                      transformers=transformers)

print(f"Processing chromosome: {ch}")

vcf_to_zarr_func(ch)

可行性说明 #

完全可以通过这种方式合并拆分的Zarr存储，但需要保证各子Zarr的结构（样本集、字段定义等）完全一致，并正确配置Zarr的组结构。以下是具体操作步骤和优化建议：

手动合并操作步骤 #

1. 创建父Zarr根目录
   先建立一个空目录作为合并后的Zarr根存储：

   mkdir merged_vcf.zarr
   

2. 迁移子Zarr内容
   将每个染色体对应的Zarr存储（如input.vcf.gz.chr1.zarr）内的所有文件，移动到父目录下以染色体命名的子目录中：

   # 示例：处理chr1
   mv input.vcf.gz.chr1.zarr/* merged_vcf.zarr/chr1/
   # 对14个染色体逐一执行此命令
   

3. 添加根组标记文件
   在merged_vcf.zarr目录下创建.zgroup文件，内容为：

   {
       "zarr_format": 2
   }
   

   该文件用于标记此目录为Zarr的根组。

4. 统一处理全局共享字段
   原代码中每个子Zarr都存储了samples字段，合并后只需保留根目录下的一份即可。从任意一个子Zarr中复制samples相关文件到父目录：

   cp input.vcf.gz.chr1.zarr/samples.* merged_vcf.zarr/
   # 可选：删除各子目录下的samples文件以减少冗余
   rm merged_vcf.zarr/chr*/samples.*
   

验证合并结果 #

用以下代码确认合并后的Zarr存储可用：

import allel
import zarr

# 打开合并后的Zarr存储
store = zarr.open('merged_vcf.zarr', mode='r')

# 检查所有染色体组是否存在
chrom_list = ['chr1', 'chr2', ...]  # 替换为你的实际染色体列表
for ch in chrom_list:
    assert ch in store, f"染色体 {ch} 未在合并存储中找到"

# 读取某条染色体的变异数据
variants = allel.VariantChunkedTable(store['chr1']['variants'])
print(f"chr1 变异数量: {len(variants)}")

# 读取样本信息
samples = store['samples'][:]
print(f"样本数量: {len(samples)}")

优化建议：修改原代码避免手动合并 #

如果后续还要生成拆分的Zarr，可以直接修改原代码，让SLURM任务写入同一个父Zarr的不同子组，省去后续手动合并的步骤：

#!/usr/bin/env python3
import sys
import os
import subprocess
import allel
import zarr

INFN = sys.argv[1]
if not INFN:
    print('Must provide input .vcf.gz as first argument')
    sys.exit(2)

FIELDS = [
    'samples',
    'variants/CHROM',
    'variants/POS',
    'variants/REF',
    'variants/ALT',
    'variants/QUAL',
    'variants/TYPE',
    'variants/is_snp',
    'variants/numalt',
    'variants/AF',
    'variants/DP',
    'variants/ANN',
    'calldata/DP',
    'calldata/GT',
]
EXCLUDE_FIELDS = None
TABIX_EXEC = 'tabix'

# 定义全局父Zarr存储路径
PARENT_ZARR = 'merged_vcf.zarr'

task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
chroms = subprocess.check_output([TABIX_EXEC,'-l',INFN],
                                 universal_newlines=True).strip().split('\n')
ch = chroms[task_id]

transformers = None
if 'ANN' in FIELDS:
    transformers=allel.ANNTransformer()

print(f"Processing chromosome: {ch}")
# 直接写入父Zarr的对应染色体子组，仅第一个任务写入samples避免重复
allel.vcf_to_zarr(INFN, PARENT_ZARR,
                  region=ch,
                  group=ch,
                  log=sys.stderr,
                  fields=FIELDS,
                  exclude_fields=EXCLUDE_FIELDS,
                  tabix=TABIX_EXEC,
                  transformers=transformers,
                  write_samples=(task_id == 0))

注意事项 #

• 确保所有子Zarr的字段定义完全一致，否则合并后读取会出现结构不匹配的错误。
• 若未删除子目录下的samples文件，读取时可能出现冲突，需保证根目录下的samples是唯一的。
• 合并过程中避免中断，防止Zarr存储损坏。

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