def split_by_key(self, key_fn: str, outfile: str, num_partitions: int, auto_mkdir: bool = False, verbose: bool = True): """Spit the partition into multiple smaller partitions by key Args: key_fn (str): function outfile (str): outfile num_partitions (int): [description] auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): [description]. Defaults to True. """ outfile = create_filepath_template(outfile, False) key_fn = get_func_by_name(key_fn) with contextlib.ExitStack() as stack, self._open() as f: writers = [ stack.enter_context( PartitionWriter(outfile.format(auto=i, stem=self.stem), auto_mkdir=auto_mkdir)) for i in range(num_partitions) ] for line in tqdm(f, total=self.n_records) if verbose else f: bucket_no = key_fn(self.deser_fn(line)) partno = bucket_no % num_partitions writers[partno].write(line)
def filter(self, fn: str, outfile: str, delete_on_empty: bool = False, auto_mkdir: bool = False, verbose: bool = True): """Filter function Args: fn (str): [description] outfile (str): [description] delete_on_empty (bool, optional): delete the partition if there is no record. Defaults to False auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): [description]. Defaults to True. """ fn = get_func_by_name(fn) outfile = create_filepath_template(outfile, True).format(auto=0, stem=self.stem) with self._open() as f, PartitionWriter( outfile, on_close_delete_if_empty=delete_on_empty, auto_mkdir=auto_mkdir) as g: for line in tqdm(f, total=self.n_records) if verbose else f: if fn(self.deser_fn(line)): g.write(line)
def flat_map(self, fn: str, outfile: str, auto_mkdir: bool = False, verbose: bool = True): """Apply a flat map function on every record of this partition Args: fn (str): an import path of the map function, which should has this signature: (record: Any) -> Any outfile (str): output file of the new partition, `*` or `{stem}` in the path are placeholders, which will be replaced by the stem (i.e., file name without extension) of the current partition auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): show the execution progress bar. Defaults to True. """ fn = get_func_by_name(fn) outfile = create_filepath_template(outfile, True).format(auto=0, stem=self.stem) with self._open() as f, PartitionWriter(outfile, auto_mkdir=auto_mkdir) as g: for record in tqdm(f, total=self.n_records) if verbose else f: record = self.deser_fn(record) for sub_record in fn(record): sub_record = self.ser_fn(sub_record) g.write(sub_record) g.write_new_line()
def reduce_by_key(self, key_fn: str, fn: str, outfile: str, init_val: Any = None, auto_mkdir: bool = False, verbose: bool = True): """Reduce Args: key_fn (str): [description] fn (str): [description] outfile (str): [description] init_val (Any): [description] auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): [description]. Defaults to True. """ key_fn = get_func_by_name(key_fn) fn = get_func_by_name(fn) outfile = create_filepath_template(outfile, True).format(auto=0, stem=self.stem) with self._open() as f, PartitionWriter(outfile, auto_mkdir=auto_mkdir) as g: groups = {} for line in tqdm(f, total=self.n_records if self.n_records is not None else None) if verbose else f: record = self.deser_fn(line) rid = key_fn(record) if rid not in groups: if init_val is None: groups[rid] = fn(record) else: groups[rid] = fn(record, init_val) else: groups[rid] = fn(record, groups[rid]) for value in tqdm( groups.values(), total=len(groups)) if verbose else groups.values(): g.write(self.ser_fn(value)) g.write_new_line()
def apply(self, fn: str, verbose: bool = True): """Apply Args: fn (str): function verbose (bool, optional): show execution progress. Defaults to True. """ fn = get_func_by_name(fn) with self._open() as f: for line in tqdm(f, total=self.n_records) if verbose else f: fn(self.deser_fn(line))
def main(args): # build parser and parse arguments start = time() parser = build_parser() args = parser.parse_args(args) end = time() args = vars(args) command = args.pop("command") verbose = args.pop("verbose") infile = args.pop("infile") skip_nrows = args.pop("skip_nrows") deser_fn = get_func_by_name(args.pop("deser_fn")) ser_fn = get_func_by_name(args.pop("ser_fn")) if verbose > 1: print(f">>> run sh map-reduce version {__version__}") print( f">>> build parser and parsing arguments takes {end - start:.03f}s" ) classname, method_name = command.rsplit(".", 1) instance = CLASSES[classname](infile, deser_fn, ser_fn, skip_nrows) fn = getattr(instance, method_name) if 'verbose' in FUNCTIONS[command]: args['verbose'] = verbose > 0 if verbose > 1: print(f">>> execute functions with arguments:\n%s" % orjson.dumps(args)) fn(**args) if verbose > 1: print(f">>> finish! execution time: {time() - start:.03f}s")
def reduce(self, fn: str, outfile: str, init_val: Any = None, auto_mkdir: bool = False, verbose: bool = True): """Reduce Args: fn (str): [description] outfile (str): [description] init_val (Any): [description] auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): [description]. Defaults to True. """ fn = get_func_by_name(fn) if init_val is not None: accum = init_val else: accum = None outfile = create_filepath_template(outfile, False).format(auto=0, stem="") with PartitionWriter(outfile, auto_mkdir=auto_mkdir) as g, (tqdm( total=self.size) if verbose else fake_tqdm()) as pbar: for inpart in self.partitions: with inpart._open() as f: if accum is None: try: record = inpart.deser_fn(next(f)) accum = fn(record) if verbose: pbar.update(1) except StopIteration: pass for line in f: record = inpart.deser_fn(line) accum = fn(record, accum) pbar.update(1) g.write(self.ser_fn(accum)) g.write_new_line()
def reduce(self, fn: str, outfile: str, init_val: Any = None, auto_mkdir: bool = False, verbose: bool = True): """Reduce Args: fn (str): [description] outfile (str): [description] init_val (Any): [description] auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): [description]. Defaults to True. """ fn = get_func_by_name(fn) outfile = create_filepath_template(outfile, True).format(auto=0, stem=self.stem) if init_val is not None: accum = init_val else: accum = None with self._open() as f, PartitionWriter(outfile, auto_mkdir=auto_mkdir) as g: if accum is None: try: record = self.deser_fn(next(f)) accum = fn(record) except StopIteration: pass for line in tqdm(f, total=self.n_records - 1 if self.n_records is not None else None) if verbose else f: record = self.deser_fn(line) accum = fn(record, accum) g.write(self.ser_fn(accum)) g.write_new_line()
def distinct(self, key_fn: str, outfile: str, auto_mkdir: bool = False, verbose: bool = True): """Create a new partition containing the distinct elements in this partition. Args: key_fn (str): key function outfile (str): output file auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): [description]. Defaults to True. """ key_fn = get_func_by_name(key_fn) outfile = create_filepath_template(outfile, True).format(auto=0, stem=self.stem) with self._open() as f, PartitionWriter(outfile, auto_mkdir=auto_mkdir) as g: keys = set() for line in tqdm(f, total=self.n_records) if verbose else f: key = key_fn(self.deser_fn(line)) if key not in keys: keys.add(key) g.write(line)
def join(self, key_fn: str, outfile: str, partition: str, partition_key_fn: str, partition_skip_nrows: int = 0, partition_deser_fn: str = None, partition_ser_fn: str = None, output_ser_fn: str = None, auto_mkdir: bool = False, verbose: bool = True): """Join two partitions together based on its key. Args: key_fn (str): function that extracts key of the current partition outfile (str): the output file partition (str): path to the other partition partition_key_fn (str): function that extracts key of the other partition partition_skip_nrows (int, optional): number of rows to skip, default to be 0 partition_deser_fn (str, optional): the deserialization function of other partition, default to be the same as current partition partition_ser_fn (str, optional): the serialization function of other partition, default to be the same as current partition output_ser_fn (str, optional): the serialization function of the output partition, default to be the same as current partition auto_mkdir (bool, optional): automatically create directory if the directory of the output file does not exist. Defaults to False verbose (bool, optional): whether printing the process bar, default to be true """ if partition_deser_fn is None: partition_deser_fn = self.deser_fn else: partition_deser_fn = get_func_by_name(partition_deser_fn) if partition_ser_fn is None: partition_ser_fn = self.ser_fn else: partition_ser_fn = get_func_by_name(partition_ser_fn) if output_ser_fn is None: output_ser_fn = self.ser_fn else: output_ser_fn = get_func_by_name(output_ser_fn) key_fn = get_func_by_name(key_fn) partition_key_fn = get_func_by_name(partition_key_fn) outfile = create_filepath_template(outfile, True).format(auto=0, stem=self.stem) other = Partition(partition, partition_deser_fn, partition_ser_fn, partition_skip_nrows) with self._open() as f, other._open() as g: join = {} used_keys = set() for line in tqdm(f, total=self.n_records, desc="process part 0") if verbose else f: record = self.deser_fn(line) key = key_fn(record) join[key] = [record] used_keys.add(key) for line in tqdm(g, total=other.n_records, desc="process part 1") if verbose else g: record = partition_deser_fn(line) key = partition_key_fn(record) if key in join: join[key].append(record) used_keys.pop(key) for key in used_keys: join.pop(key) with PartitionWriter(outfile, auto_mkdir=auto_mkdir) as k: for r in tqdm( join.values(), total=len(join), desc="writing result") if verbose else join.values(): k.write(output_ser_fn(r)) k.write_new_line()