Nest multiple variable inputs into a single record or list of batch records.

    Custom CWL implementations extract and merge these.
496693
def _nest_vars_in_rec(var_items, rec_items, input_order, items_by_key, parallel):
    
    num_items = var_items
    var_items = list(var_items)[0]
    if rec_items:
        rec_items = list(rec_items)[0]
        if ((rec_items == 1 and var_items > 1) or parallel.startswith("batch")):
            num_items = set([rec_items])
            for var_key in (k for (k, t) in input_order.items() if t != "record"):
                var_key = tuple(var_key.split("__"))
                items_by_key[var_key] = [items_by_key[var_key]] * rec_items
        else:
            assert var_items == rec_items, (var_items, rec_items)
    return items_by_key, num_items
582938
def cbow_batch(centers, contexts, num_tokens, dtype, index_dtype):
    """"""
    contexts_data, contexts_row, contexts_col = contexts
    centers = mx.nd.array(centers, dtype=index_dtype)
    contexts = mx.nd.sparse.csr_matrix(
        (contexts_data, (contexts_row, contexts_col)),
        dtype=dtype, shape=(len(centers), num_tokens))  # yapf: disable
    return centers, contexts
391327
def _collapse_to_cwl_record(samples, want_attrs, input_files):
    
    input_keys = sorted(list(set().union(*[d["cwl_keys"] for d in samples])), key=lambda x: (-len(x), tuple(x)))
    out = {}
    for key in input_keys:
        if key in want_attrs:
            key_parts = key.split("__")
            vals = []
            cur = []
            for d in samples:
                vals.append(_to_cwl(tz.get_in(key_parts, d), input_files))
                # Remove nested keys to avoid specifying multiple times
                cur.append(_dissoc_in(d, key_parts) if len(key_parts) > 1 else d)
            samples = cur
            out[key] = vals
    return out
358297
def _read_from_cwlinput(in_file, work_dir, runtime, parallel, input_order, output_cwl_keys):
    
    with open(in_file) as in_handle:
        inputs = json.load(in_handle)
    items_by_key = {}
    input_files = []
    passed_keys = set([])
    for key, input_val in ((k, v) for (k, v) in inputs.items() if not k.startswith(("sentinel", "ignore"))):
        if key.endswith("_toolinput"):
            key = key.replace("_toolinput", "")
        if input_order[key] == "record":
            cur_keys, items = _read_cwl_record(input_val)
            passed_keys |= cur_keys
            items_by_key[key] = items
        else:
            items_by_key[tuple(key.split("__"))] = _cwlvar_to_wdl(input_val)
        input_files = _find_input_files(input_val, input_files)
    prepped = _merge_cwlinputs(items_by_key, input_order, parallel)
    out = []
    for data in prepped:
        if isinstance(data, (list, tuple)):
            out.append([_finalize_cwl_in(utils.to_single_data(x), work_dir, list(passed_keys),
                                         output_cwl_keys, runtime) for x in data])
        else:
            out.append(_finalize_cwl_in(data, work_dir, list(passed_keys), output_cwl_keys, runtime))
    return out, input_files
474667
def _expand_rec_to_vars(var_items, rec_items, input_order, items_by_key, parallel):
    
    num_items = var_items
    var_items = list(var_items)[0]
    if rec_items:
        for rec_key in (k for (k, t) in input_order.items() if t == "record"):
            rec_vals = items_by_key[rec_key]
            if len(rec_vals) == 1 and var_items > 1:
                items_by_key[rec_key] = rec_vals * var_items
            else:
                assert var_items == len(rec_vals), (var_items, rec_vals)
    return items_by_key, num_items
666302
def _batch_insert(self, inserts, intervals, **kwargs):
    '''
    
    '''
    for timestamp,names in inserts.iteritems():
      for name,values in names.iteritems():
        for value in values:
          self._insert( name, value, timestamp, intervals, **kwargs )
1334443
def _batch_to_stdin(process, in_record_sep, in_column_sep, in_batches, batch_to_file_s):
        
        for i, b2f in enumerate(batch_to_file_s):
            if b2f.is_stdin():
                input_str = BaseShellOperator._input_str(in_batches[i], in_record_sep, in_column_sep)
                b2f.write_stdin(process.stdin, input_str)
                break
207636
def summarize_vc(items):
    
    items = [utils.to_single_data(x) for x in utils.flatten(items)]
    items = [_normalize_vc_input(x) for x in items]
    items = validate.summarize_grading(items)
    items = [utils.to_single_data(x) for x in items]
    out = {"validate": validate.combine_validations(items),
           "variants": {"calls": [], "gvcf": [], "samples": []}}
    added = set([])
    variants_by_sample = collections.defaultdict(list)
    sample_order = []
    for data in items:
        batch_samples = data.get("batch_samples", [dd.get_sample_name(data)])
        for s in batch_samples:
            if s not in sample_order:
                sample_order.append(s)
        if data.get("vrn_file"):
            # Only get batches if we're actually doing variantcalling in bcbio
            # otherwise we'll be using the original files
            names = dd.get_batches(data) if dd.get_variantcaller(data) else None
            if not names:
                names = [dd.get_sample_name(data)]
            batch_name = names[0]
            if data.get("vrn_file_joint") is not None:
                to_add = [("vrn_file", "gvcf", dd.get_sample_name(data)),
                          ("vrn_file_joint", "calls", batch_name)]
            else:
                to_add = [("vrn_file", "calls", batch_name)]
            for vrn_key, out_key, name in to_add:
                cur_name = "%s-%s" % (name, dd.get_variantcaller(data))
                out_file = os.path.join(utils.safe_makedir(os.path.join(dd.get_work_dir(data),
                                                                        "variants", out_key)),
                                        "%s.vcf.gz" % cur_name)
                for s in batch_samples:
                    variants_by_sample[s].append(out_file)
                if cur_name not in added:
                    added.add(cur_name)
                    # Ideally could symlink here but doesn't appear to work with
                    # Docker container runs on Toil where PATHs don't get remapped
                    utils.copy_plus(os.path.realpath(data[vrn_key]), out_file)
                    vcfutils.bgzip_and_index(out_file, data["config"])
                    out["variants"][out_key].append(out_file)
    for sample in sample_order:
        out["variants"]["samples"].append(variants_by_sample[sample])
    return [out]
810174
def batch_iterable(l, n):
    ''' '''
    i = iter(l)
    piece = list(islice(i, n))
    while piece:
        yield piece
        piece = list(islice(i, n))
917051
def _normalize_cwl_inputs(items):
    
    with_validate = {}
    vrn_files = []
    ready_items = []
    batch_samples = []
    for data in (cwlutils.normalize_missing(utils.to_single_data(d)) for d in items):
        batch_samples.append(dd.get_sample_name(data))
        if tz.get_in(["config", "algorithm", "validate"], data):
            with_validate[_checksum(tz.get_in(["config", "algorithm", "validate"], data))] = data
        if data.get("vrn_file"):
            vrn_files.append(data["vrn_file"])
        ready_items.append(data)
    if len(with_validate) == 0:
        data = _pick_lead_item(ready_items)
        data["batch_samples"] = batch_samples
        return data
    else:
        assert len(with_validate) == 1, len(with_validate)
        assert len(set(vrn_files)) == 1, set(vrn_files)
        data = _pick_lead_item(with_validate.values())
        data["batch_samples"] = batch_samples
        data["vrn_file"] = vrn_files[0]
        return data
44201
def prep_cwl(samples, workflow_fn, out_dir, out_file, integrations=None,
             add_container_tag=None):
    
    if add_container_tag is None:
        container_tags = None
    elif add_container_tag.lower() == "quay_lookup":
        container_tags = {}
    else:
        container_tags = collections.defaultdict(lambda: add_container_tag)
    step_dir = utils.safe_makedir(os.path.join(out_dir, "steps"))
    get_retriever = GetRetriever(integrations, samples)
    variables, keyvals = _flatten_samples(samples, out_file, get_retriever)
    cur_remotes = _get_cur_remotes(keyvals)
    file_estimates = _calc_input_estimates(keyvals, get_retriever)
    out = _cwl_workflow_template(variables)
    parent_wfs = []
    step_parallelism = {}
    steps, wfoutputs = workflow_fn(samples)
    used_inputs = set([])
    for cur in workflow.generate(variables, steps, wfoutputs):
        if cur[0] == "step":
            _, name, parallel, inputs, outputs, image, programs, disk, cores, no_files = cur
            step_file = _write_tool(step_dir, name, inputs, outputs, parallel, image, programs,
                                    file_estimates, disk, cores, samples, cur_remotes, no_files, container_tags)
            out["steps"].append(_step_template(name, step_file, inputs, outputs, parallel, step_parallelism))
            used_inputs |= set(x["id"] for x in inputs)
        elif cur[0] == "expressiontool":
            _, name, inputs, outputs, expression, parallel = cur
            step_file = _write_expressiontool(step_dir, name, inputs, outputs, expression, parallel)
            out["steps"].append(_step_template(name, step_file, inputs, outputs, parallel, step_parallelism))
            used_inputs |= set(x["id"] for x in inputs)
        elif cur[0] == "upload":
            for output in cur[1]:
                wf_output = copy.deepcopy(output)
                if "outputSource" not in wf_output:
                    wf_output["outputSource"] = wf_output.pop("source")
                wf_output = _clean_record(wf_output)
                # Avoid input/output naming clashes
                if wf_output["id"] in used_inputs:
                    wf_output["id"] = "%s_out" % wf_output["id"]
                out["outputs"].append(wf_output)
        elif cur[0] == "wf_start":
            parent_wfs.append(out)
            out = _cwl_workflow_template(cur[1])
        elif cur[0] == "wf_finish":
            _, name, parallel, inputs, outputs, scatter = cur
            wf_out_file = "wf-%s.cwl" % name
            with open(os.path.join(out_dir, wf_out_file), "w") as out_handle:
                yaml.safe_dump(out, out_handle, default_flow_style=False, allow_unicode=False)
            out = parent_wfs.pop(-1)
            out["steps"].append(_step_template(name, wf_out_file, inputs, outputs, parallel,
                                               step_parallelism, scatter))
            used_inputs |= set(x["id"] for x in inputs)
        else:
            raise ValueError("Unexpected workflow value %s" % str(cur))
        step_parallelism[name] = parallel

    with open(out_file, "w") as out_handle:
        out["inputs"] = [x for x in out["inputs"] if x["id"] in used_inputs]
        yaml.safe_dump(out, out_handle, default_flow_style=False, allow_unicode=False)
    sample_json = "%s-samples.json" % utils.splitext_plus(out_file)[0]
    out_clean = _clean_final_outputs(copy.deepcopy({k: v for k, v in keyvals.items() if k in used_inputs}),
                                     get_retriever)
    with open(sample_json, "w") as out_handle:
        json.dump(out_clean, out_handle, sort_keys=True, indent=4, separators=(',', ': '))
    return out_file, sample_json
514997
def _do_scatter_var(v, parallel):
    
    # For batches, scatter records only at the top level (double nested)
    if parallel.startswith("batch") and workflow.is_cwl_record(v):
        return (tz.get_in(["type", "type"], v) == "array" and
                tz.get_in(["type", "type", "type"], v) == "array")
    # Otherwise, scatter arrays
    else:
        return (tz.get_in(["type", "type"], v) == "array")
1203404
def group_batches_joint(samples):
    
    def _caller_batches(data):
        jointcaller = tz.get_in(("config", "algorithm", "jointcaller"), data)
        batch = tz.get_in(("metadata", "batch"), data) if jointcaller else None
        return jointcaller, batch
    def _prep_data(data, items):
        for r in ["callable_regions", "variant_regions"]:
            data[r] = list(set(filter(lambda x: x is not None,
                                      [tz.get_in(("config", "algorithm", r), d) for d in items])))
        data["work_bams"] = [dd.get_align_bam(x) or dd.get_work_bam(x) for x in items]
        data["vrn_files"] = [x["vrn_file"] for x in items]
        return data
    return _group_batches_shared(samples, _caller_batches, _prep_data)
1095222
def concatenate_variables(scope, variables, container):
    '''
    
    '''

    # Check if it's possible to concatenate those inputs.
    type_set = set(type(variable.type) for variable in variables)
    number_type_set = {FloatType, FloatTensorType, Int64Type, Int64TensorType}
    if StringType in type_set and any(number_type in type_set for number_type in number_type_set):
        raise RuntimeError('We are not able to concatenate numerical tensor(s) and string tensor(s)')

    input_names = []  # input variables' names we want to concatenate
    input_dims = []  # dimensions of the variables that is going to be concatenated

    # Collect input variable names and do cast if needed
    for variable in variables:
        if isinstance(variable.type, (Int64TensorType, Int64Type)):
            input_names.append(convert_integer_to_float(scope, variable, container))
        else:
            input_names.append(variable.full_name)
        # We assume input variables' shape are [1, C_1], ..., [1, C_n] if there are n inputs.
        input_dims.append(variable.type.shape[1])

    if len(input_names) == 1:
        # No need to concatenate tensors if there is only one input
        return input_names[0]
    else:
        # To combine all inputs, we need a FeatureVectorizer
        op_type = 'FeatureVectorizer'
        attrs = {'name': scope.get_unique_operator_name(op_type), 'inputdimensions': input_dims}
        # Create a variable name to capture feature vectorizer's output
        concatenated_name = scope.get_unique_variable_name('concatenated')
        # Set up our FeatureVectorizer
        container.add_node(op_type, input_names, concatenated_name, op_domain='ai.onnx.ml', **attrs)

        return concatenated_name
1011655
def _batch_insert(self, inserts, intervals, **kwargs):
    '''
    
    '''
    if 'pipeline' in kwargs:
      pipe = kwargs.get('pipeline')
      own_pipe = False
    else:
      pipe = self._client.pipeline(transaction=False)
      kwargs['pipeline'] = pipe
      own_pipe = True

    ttl_batch = set()
    for timestamp,names in inserts.iteritems():
      for name,values in names.iteritems():
        for value in values:
          # TODO: support config param to flush the pipe every X inserts
          self._insert( name, value, timestamp, intervals, ttl_batch=ttl_batch, **kwargs )

    for ttl_args in ttl_batch:
      pipe.expire(*ttl_args)

    if own_pipe:
      kwargs['pipeline'].execute()
1188636
def _create_record(name, field_defs, step_name, inputs, unlist, file_vs, std_vs, parallel):
    
    if field_defs:
        fields = []
        inherit = []
        inherit_all = False
        inherit_exclude = []
        for fdef in field_defs:
            if not fdef.get("type"):
                if fdef["id"] == "inherit":
                    inherit_all = True
                    inherit_exclude = fdef.get("exclude", [])
                else:
                    inherit.append(fdef["id"])
            else:
                cur = {"name": _get_string_vid(fdef["id"]),
                       "type": fdef["type"]}
                fields.append(_add_secondary_to_rec_field(fdef, cur))
        if inherit_all:
            fields.extend(_infer_record_outputs(inputs, unlist, file_vs, std_vs, parallel, exclude=inherit_exclude))
        elif inherit:
            fields.extend(_infer_record_outputs(inputs, unlist, file_vs, std_vs, parallel, inherit))
    else:
        fields = _infer_record_outputs(inputs, unlist, file_vs, std_vs, parallel)
    out = {"id": "%s/%s" % (step_name, name),
           "type": {"name": name,
                    "type": "record",
                    "fields": fields}}
    if parallel in ["batch-single", "multi-batch"]:
        out = _nest_variable(out)
    return out
760923
def create_document_batches(jsonlines, id_field, max_batch_size=CLOUDSEARCH_MAX_BATCH_SIZE):
    
    batch = []
    fixed_initial_size = 2

    def create_entry(line):
        try:
            record = json.loads(line)
        except:
            raise ValueError('Could not parse JSON from: %s' % line)
        key = record[id_field]
        return '{"type":"add","id":%s,"fields":%s}' % (json.dumps(key), line)

    current_size = fixed_initial_size
    for line in jsonlines:
        entry = create_entry(line)
        entry_size = len(entry) + 1

        if max_batch_size > (current_size + entry_size):
            current_size += entry_size
            batch.append(entry)
        else:
            yield '[' + ','.join(batch) + ']'
            batch = [entry]
            current_size = fixed_initial_size + entry_size

    if batch:
        yield '[' + ','.join(batch) + ']'