collective/cdsc_reddit
18
0
Fork 0
Code Activity

datasets/helper.py: use zstandard library for .zst decompression

Browse Source 
The Python environment runs inside a Singularity container that cannot
exec the host's /usr/bin/zstd via subprocess. Replace the subprocess
call with the zstandard Python library, which was already a dependency.
Other formats (bz2, xz, gz) still use subprocess as before.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Benjamin Mako Hill
2026-05-25 18:48:04 -07:00
parent 18925dfe5b
commit bf6ccbc84a
1 changed files with 20 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
datasets/helper.py
View File

@@ -3,6 +3,9 @@ import re
from collections import defaultdict from collections import defaultdict
from os import path from os import path
import glob import glob
import io
import zstandard
def find_dumps(dumpdir, base_pattern): def find_dumps(dumpdir, base_pattern):
@@ -28,24 +31,28 @@ def open_fileset(files):
yield line yield line
def open_input_file(input_filename): def open_input_file(input_filename):
# .zst handled via the zstandard library to avoid subprocess/container issues
if re.match(r'.*\.zst$', input_filename):
fh = open(input_filename, 'rb')
dctx = zstandard.ZstdDecompressor()
return io.TextIOWrapper(dctx.stream_reader(fh), encoding='utf-8')
if re.match(r'.*\.7z$', input_filename): if re.match(r'.*\.7z$', input_filename):
cmd = ["7za", "x", "-so", input_filename, '*'] cmd = ["7za", "x", "-so", input_filename, '*']
elif re.match(r'.*\.gz$', input_filename):
cmd = ["zcat", input_filename]
elif re.match(r'.*\.bz2$', input_filename): elif re.match(r'.*\.bz2$', input_filename):
cmd = ["bzcat", "-dk", input_filename] cmd = ["bzcat", "-dk", input_filename]
elif re.match(r'.*\.bz', input_filename): elif re.match(r'.*\.bz', input_filename):
cmd = ["bzcat", "-dk", input_filename] cmd = ["bzcat", "-dk", input_filename]
elif re.match(r'.*\.xz', input_filename): elif re.match(r'.*\.xz', input_filename):
cmd = ["xzcat", '-dk', '-T 20', input_filename] cmd = ["xzcat", '-dk', '-T 20', input_filename]
elif re.match(r'.*\.zst',input_filename):
cmd = ['zstd','-dck', input_filename]
elif re.match(r'.*\.gz', input_filename): elif re.match(r'.*\.gz', input_filename):
cmd = ['gzip','-dc', input_filename] cmd = ["zcat", input_filename]
else:
return open(input_filename, 'r')
try: try:
input_file = Popen(cmd, stdout=PIPE).stdout return Popen(cmd, stdout=PIPE).stdout
except NameError as e: except NameError as e:
print(e) print(e)
input_file = open(input_filename, 'r') return open(input_filename, 'r')
return input_file