Function get_merging_graph

ciphercore-base/src/mpc/mpc_psi.rs:407–453

Convert key columns to binary and merge them for each input database

(
    context: Context,
    column_types: ColumnTypesInfo,
    key_headers: &[String],
    is_private: bool,
)

// Convert key columns to binary and merge them for each input database
fn get_merging_graph(
    context: Context,
    column_types: ColumnTypesInfo,
    key_headers: &[String],
    is_private: bool,
) -> Result<Graph> {
    let mut headers_map = HashMap::new();
    for (h, t) in &column_types.headers_types {
        headers_map.insert((*h).clone(), (*t).clone());
    }

    let merging_context = simple_context(|g| {
        let data = g.input(column_types.get_plaintext_type())?;

        let num_entries = column_types.get_num_entries();

        let mut bit_columns = vec![];
        for header in key_headers {
            let t = headers_map.get(header).unwrap();

            let mut column = data.named_tuple_get((*header).clone())?;
            if t.has_mask() {
                column = column.tuple_get(1)?;
            }
            let mut bit_column = if t.get_scalar_type() != BIT {
                column.a2b()?
            } else {
                column
            };
            // Flatten all the bits per entry
            let flattened_shape = vec![num_entries, t.get_row_size_in_bits()];
            bit_column = bit_column.reshape(array_type(flattened_shape, BIT))?;
            // Pull out bits to simplify merging of columns
            bit_columns.push(bit_column);
        }
        // Merge key columns
        let merged_columns = if bit_columns.len() > 1 {
            g.concatenate(bit_columns, 1)?
        } else {
            bit_columns[0].clone()
        };

        Ok(merged_columns)
    })?;

    convert_main_graph_to_mpc(merging_context, context, vec![is_private])
}

// Multiply the column by the mask/null column
// Assumes that column mask is binary and it has the same number of rows as the column