mirror of
https://github.com/galaxyproject/galaxy.git
synced 2026-09-24 16:30:27 +08:00
451 lines
16 KiB
Python
451 lines
16 KiB
Python
#!/usr/bin/env python
|
|
# Dan Blankenberg
|
|
"""
|
|
Script to Join Two Files on specified columns.
|
|
|
|
Takes two tab delimited files, two column numbers (base 1) and outputs a new tab delimited file with lines joined by tabs.
|
|
User can also opt to have have non-joining rows of file1 echoed.
|
|
"""
|
|
from __future__ import print_function
|
|
|
|
import json
|
|
import optparse
|
|
import os
|
|
import re
|
|
import struct
|
|
import sys
|
|
import tempfile
|
|
|
|
from galaxy.util import stringify_dictionary_keys
|
|
from galaxy.util.bunch import Bunch
|
|
|
|
|
|
class OffsetList(object):
|
|
def __init__(self, filesize=0, fmt=None):
|
|
self.file = tempfile.NamedTemporaryFile("w+b")
|
|
if fmt:
|
|
self.fmt = fmt
|
|
elif filesize and filesize <= sys.maxsize * 2:
|
|
self.fmt = "I"
|
|
else:
|
|
self.fmt = "Q"
|
|
self.fmt_size = struct.calcsize(self.fmt)
|
|
|
|
@property
|
|
def size(self):
|
|
self.file.flush()
|
|
return self.file_size / self.fmt_size
|
|
|
|
@property
|
|
def file_size(self):
|
|
self.file.flush()
|
|
return os.stat(self.file.name).st_size
|
|
|
|
def add_offset(self, offset):
|
|
if not isinstance(offset, list):
|
|
offset = [offset]
|
|
self.file.seek(self.file_size)
|
|
for off in offset:
|
|
self.file.write(struct.pack(self.fmt, off))
|
|
|
|
def get_offsets(self, start=0):
|
|
self.file.seek(start * self.fmt_size)
|
|
while True:
|
|
packed = self.file.read(self.fmt_size)
|
|
if not packed:
|
|
break
|
|
yield struct.unpack(self.fmt, packed)[0]
|
|
|
|
def get_offset_by_index(self, index):
|
|
self.file.seek(index * self.fmt_size)
|
|
return struct.unpack(self.fmt, self.file.read(self.fmt_size))[0]
|
|
|
|
def set_offset_at_index(self, index, offset):
|
|
if not isinstance(offset, list):
|
|
offset = [offset]
|
|
if index >= self.size:
|
|
self.add_offset(offset)
|
|
else:
|
|
temp_file = tempfile.NamedTemporaryFile("w+b")
|
|
self.file.seek(0)
|
|
temp_file.write(self.file.read((index) * self.fmt_size))
|
|
for off in offset:
|
|
temp_file.write(struct.pack(self.fmt, off))
|
|
temp_file.write(self.file.read())
|
|
self.file = temp_file
|
|
|
|
|
|
class SortedOffsets(OffsetList):
|
|
def __init__(self, indexed_filename, column, split=None):
|
|
OffsetList.__init__(self, os.stat(indexed_filename).st_size)
|
|
self.indexed_filename = indexed_filename
|
|
self.indexed_file = open(indexed_filename)
|
|
self.column = column
|
|
self.split = split
|
|
self.last_identifier = None
|
|
self.last_identifier_merged = None
|
|
self.last_offset_merged = 0
|
|
|
|
def merge_with_dict(self, new_offset_dict):
|
|
if not new_offset_dict:
|
|
return # no items to merge in
|
|
keys = list(new_offset_dict.keys())
|
|
keys.sort()
|
|
identifier2 = keys.pop(0)
|
|
|
|
result_offsets = OffsetList(fmt=self.fmt)
|
|
offsets1 = enumerate(self.get_offsets())
|
|
try:
|
|
index1, offset1 = next(offsets1)
|
|
identifier1 = self.get_identifier_by_offset(offset1)
|
|
except StopIteration:
|
|
offset1 = None
|
|
identifier1 = None
|
|
index1 = 0
|
|
|
|
while True:
|
|
if identifier1 is None and identifier2 is None:
|
|
self.file = result_offsets.file # self is now merged results
|
|
return
|
|
elif identifier1 is None or (identifier2 and identifier2 < identifier1):
|
|
result_offsets.add_offset(new_offset_dict[identifier2])
|
|
if keys:
|
|
identifier2 = keys.pop(0)
|
|
else:
|
|
identifier2 = None
|
|
elif identifier2 is None:
|
|
result_offsets.file.seek(result_offsets.file_size)
|
|
self.file.seek(index1 * self.fmt_size)
|
|
result_offsets.file.write(self.file.read())
|
|
identifier1 = None
|
|
offset1 = None
|
|
else:
|
|
result_offsets.add_offset(offset1)
|
|
try:
|
|
index1, offset1 = next(offsets1)
|
|
identifier1 = self.get_identifier_by_offset(offset1)
|
|
except StopIteration:
|
|
offset1 = None
|
|
identifier1 = None
|
|
index1 += 1
|
|
|
|
# methods to help link offsets to lines, ids, etc
|
|
def get_identifier_by_line(self, line):
|
|
if isinstance(line, str):
|
|
fields = line.rstrip("\r\n").split(self.split)
|
|
if self.column < len(fields):
|
|
return fields[self.column]
|
|
return None
|
|
|
|
def get_line_by_offset(self, offset):
|
|
self.indexed_file.seek(offset)
|
|
return self.indexed_file.readline()
|
|
|
|
def get_identifier_by_offset(self, offset):
|
|
return self.get_identifier_by_line(self.get_line_by_offset(offset))
|
|
|
|
|
|
# indexed set of offsets, index is built on demand
|
|
class OffsetIndex(object):
|
|
def __init__(self, filename, column, split=None, index_depth=3):
|
|
self.filename = filename
|
|
self.file = open(filename)
|
|
self.column = column
|
|
self.split = split
|
|
self._offsets = {}
|
|
self._index = None
|
|
self.index_depth = index_depth
|
|
|
|
def _build_index(self):
|
|
self._index = {}
|
|
for start_char, sorted_offsets in self._offsets.items():
|
|
self._index[start_char] = {}
|
|
for i, offset in enumerate(sorted_offsets.get_offsets()):
|
|
identifier = sorted_offsets.get_identifier_by_offset(offset)
|
|
if identifier[0 : self.index_depth] not in self._index[start_char]:
|
|
self._index[start_char][identifier[0 : self.index_depth]] = i
|
|
|
|
def get_lines_by_identifier(self, identifier):
|
|
if not identifier:
|
|
return
|
|
# if index doesn't exist, build it
|
|
if self._index is None:
|
|
self._build_index()
|
|
|
|
# identifier cannot exist
|
|
if identifier[0] not in self._index or identifier[0 : self.index_depth] not in self._index[identifier[0]]:
|
|
return
|
|
# identifier might exist, search for it
|
|
offset_index = self._index[identifier[0]][identifier[0 : self.index_depth]]
|
|
while True:
|
|
if offset_index >= self._offsets[identifier[0]].size:
|
|
return
|
|
offset = self._offsets[identifier[0]].get_offset_by_index(offset_index)
|
|
identifier2 = self._offsets[identifier[0]].get_identifier_by_offset(offset)
|
|
if not identifier2 or identifier2 > identifier:
|
|
return
|
|
if identifier2 == identifier:
|
|
yield self._offsets[identifier[0]].get_line_by_offset(offset)
|
|
offset_index += 1
|
|
|
|
def get_offsets(self):
|
|
keys = list(self._offsets.keys())
|
|
keys.sort()
|
|
for key in keys:
|
|
for offset in self._offsets[key].get_offsets():
|
|
yield offset
|
|
|
|
def get_line_by_offset(self, offset):
|
|
self.file.seek(offset)
|
|
return self.file.readline()
|
|
|
|
def get_identifiers_offsets(self):
|
|
keys = list(self._offsets.keys())
|
|
keys.sort()
|
|
for key in keys:
|
|
for offset in self._offsets[key].get_offsets():
|
|
yield self._offsets[key].get_identifier_by_offset(offset), offset
|
|
|
|
def get_identifier_by_line(self, line):
|
|
if isinstance(line, str):
|
|
fields = line.rstrip("\r\n").split(self.split)
|
|
if self.column < len(fields):
|
|
return fields[self.column]
|
|
return None
|
|
|
|
def merge_with_dict(self, d):
|
|
if not d:
|
|
return # no data to merge
|
|
self._index = None
|
|
keys = list(d.keys())
|
|
keys.sort()
|
|
identifier = keys.pop(0)
|
|
first_char = identifier[0]
|
|
temp = {identifier: d[identifier]}
|
|
while True:
|
|
if not keys:
|
|
if first_char not in self._offsets:
|
|
self._offsets[first_char] = SortedOffsets(self.filename, self.column, self.split)
|
|
self._offsets[first_char].merge_with_dict(temp)
|
|
return
|
|
identifier = keys.pop(0)
|
|
if identifier[0] == first_char:
|
|
temp[identifier] = d[identifier]
|
|
else:
|
|
if first_char not in self._offsets:
|
|
self._offsets[first_char] = SortedOffsets(self.filename, self.column, self.split)
|
|
self._offsets[first_char].merge_with_dict(temp)
|
|
temp = {identifier: d[identifier]}
|
|
first_char = identifier[0]
|
|
|
|
|
|
class BufferedIndex(object):
|
|
def __init__(self, filename, column, split=None, buffer=1000000, index_depth=3):
|
|
self.index = OffsetIndex(filename, column, split, index_depth)
|
|
self.buffered_offsets = {}
|
|
f = open(filename)
|
|
offset = f.tell()
|
|
identified_offset_count = 1
|
|
while True:
|
|
offset = f.tell()
|
|
line = f.readline()
|
|
if not line:
|
|
break # EOF
|
|
identifier = self.index.get_identifier_by_line(line)
|
|
if identifier:
|
|
# flush buffered offsets, if buffer size reached
|
|
if buffer and identified_offset_count % buffer == 0:
|
|
self.index.merge_with_dict(self.buffered_offsets)
|
|
self.buffered_offsets = {}
|
|
if identifier not in self.buffered_offsets:
|
|
self.buffered_offsets[identifier] = []
|
|
self.buffered_offsets[identifier].append(offset)
|
|
identified_offset_count += 1
|
|
f.close()
|
|
|
|
def get_lines_by_identifier(self, identifier):
|
|
for line in self.index.get_lines_by_identifier(identifier):
|
|
yield line
|
|
if identifier in self.buffered_offsets:
|
|
for offset in self.buffered_offsets[identifier]:
|
|
yield self.index.get_line_by_offset(offset)
|
|
|
|
|
|
def fill_empty_columns(line, split, fill_values):
|
|
if not fill_values:
|
|
return line
|
|
filled_columns = []
|
|
for i, field in enumerate(line.split(split)):
|
|
if field or i >= len(fill_values):
|
|
filled_columns.append(field)
|
|
else:
|
|
filled_columns.append(fill_values[i])
|
|
if len(fill_values) > len(filled_columns):
|
|
filled_columns.extend(fill_values[len(filled_columns) :])
|
|
return split.join(filled_columns)
|
|
|
|
|
|
def join_files(
|
|
filename1,
|
|
column1,
|
|
filename2,
|
|
column2,
|
|
out_filename,
|
|
split=None,
|
|
buffer=1000000,
|
|
keep_unmatched=False,
|
|
keep_partial=False,
|
|
keep_headers=False,
|
|
index_depth=3,
|
|
fill_options=None,
|
|
):
|
|
# return identifier based upon line
|
|
def get_identifier_by_line(line, column, split=None):
|
|
if isinstance(line, str):
|
|
fields = line.rstrip("\r\n").split(split)
|
|
if column < len(fields):
|
|
return fields[column]
|
|
return None
|
|
|
|
if fill_options is None:
|
|
fill_options = Bunch(fill_unjoined_only=True, file1_columns=None, file2_columns=None)
|
|
keep_headers_done = False
|
|
out = open(out_filename, "w")
|
|
index = BufferedIndex(filename2, column2, split, buffer, index_depth)
|
|
for line1 in open(filename1):
|
|
if keep_headers and not keep_headers_done:
|
|
header1 = line1
|
|
with open(filename2) as file2:
|
|
header2 = file2.readline()
|
|
header2 = re.sub(r"^#", "", header2)
|
|
out.write("%s%s%s\n" % (header1.rstrip("\r\n"), split, header2.rstrip("\r\n")))
|
|
keep_headers_done = True
|
|
continue
|
|
identifier = get_identifier_by_line(line1, column1, split)
|
|
if identifier:
|
|
written = False
|
|
for line2 in index.get_lines_by_identifier(identifier):
|
|
if not fill_options.fill_unjoined_only:
|
|
out.write(
|
|
"%s%s%s\n"
|
|
% (
|
|
fill_empty_columns(line1.rstrip("\r\n"), split, fill_options.file1_columns),
|
|
split,
|
|
fill_empty_columns(line2.rstrip("\r\n"), split, fill_options.file2_columns),
|
|
)
|
|
)
|
|
else:
|
|
out.write("%s%s%s\n" % (line1.rstrip("\r\n"), split, line2.rstrip("\r\n")))
|
|
written = True
|
|
if not written and keep_unmatched:
|
|
out.write(fill_empty_columns(line1.rstrip("\r\n"), split, fill_options.file1_columns))
|
|
if fill_options:
|
|
if fill_options.file2_columns:
|
|
out.write("%s%s" % (split, fill_empty_columns("", split, fill_options.file2_columns)))
|
|
out.write("\n")
|
|
elif keep_partial:
|
|
out.write(fill_empty_columns(line1.rstrip("\r\n"), split, fill_options.file1_columns))
|
|
if fill_options:
|
|
if fill_options.file2_columns:
|
|
out.write("%s%s" % (split, fill_empty_columns("", split, fill_options.file2_columns)))
|
|
out.write("\n")
|
|
out.close()
|
|
|
|
|
|
def main():
|
|
parser = optparse.OptionParser()
|
|
parser.add_option(
|
|
"-b",
|
|
"--buffer",
|
|
dest="buffer",
|
|
type="int",
|
|
default=1000000,
|
|
help="Number of lines to buffer at a time. Default: 1,000,000 lines. A buffer of 0 will attempt to use memory only.",
|
|
)
|
|
parser.add_option(
|
|
"-d",
|
|
"--index_depth",
|
|
dest="index_depth",
|
|
type="int",
|
|
default=3,
|
|
help="Depth to use on filebased offset indexing. Default: 3.",
|
|
)
|
|
parser.add_option(
|
|
"-p",
|
|
"--keep_partial",
|
|
action="store_true",
|
|
dest="keep_partial",
|
|
default=False,
|
|
help="Keep rows in first input which are missing identifiers.",
|
|
)
|
|
parser.add_option(
|
|
"-u",
|
|
"--keep_unmatched",
|
|
action="store_true",
|
|
dest="keep_unmatched",
|
|
default=False,
|
|
help="Keep rows in first input which are not joined with the second input.",
|
|
)
|
|
parser.add_option(
|
|
"-f",
|
|
"--fill_options_file",
|
|
dest="fill_options_file",
|
|
type="str",
|
|
default=None,
|
|
help="Fill empty columns with a values from a JSONified file.",
|
|
)
|
|
parser.add_option(
|
|
"-H", "--keep_headers", action="store_true", dest="keep_headers", default=False, help="Keep the headers"
|
|
)
|
|
|
|
options, args = parser.parse_args()
|
|
|
|
fill_options = None
|
|
if options.fill_options_file is not None:
|
|
try:
|
|
fill_options = Bunch(
|
|
**stringify_dictionary_keys(json.load(open(options.fill_options_file)))
|
|
) # json.load( open( options.fill_options_file ) )
|
|
except Exception as e:
|
|
print("Warning: Ignoring fill options due to json error (%s)." % e)
|
|
if fill_options is None:
|
|
fill_options = Bunch()
|
|
if "fill_unjoined_only" not in fill_options:
|
|
fill_options.fill_unjoined_only = True
|
|
if "file1_columns" not in fill_options:
|
|
fill_options.file1_columns = None
|
|
if "file2_columns" not in fill_options:
|
|
fill_options.file2_columns = None
|
|
|
|
try:
|
|
filename1 = args[0]
|
|
filename2 = args[1]
|
|
column1 = int(args[2]) - 1
|
|
column2 = int(args[3]) - 1
|
|
out_filename = args[4]
|
|
except Exception:
|
|
print("Error parsing command line.", file=sys.stderr)
|
|
sys.exit()
|
|
|
|
# Character for splitting fields and joining lines
|
|
split = "\t"
|
|
|
|
return join_files(
|
|
filename1,
|
|
column1,
|
|
filename2,
|
|
column2,
|
|
out_filename,
|
|
split,
|
|
options.buffer,
|
|
options.keep_unmatched,
|
|
options.keep_partial,
|
|
options.keep_headers,
|
|
options.index_depth,
|
|
fill_options=fill_options,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|