Class QueryEngine

union/D3L/d3l/querying/query_engine.py:12–244

class QueryEngine:
    def __init__(self, *query_backends: SimilarityIndex):
        """
        Create a new querying engine to perform data discovery in datalakes.
        Parameters
        ----------
        query_backends : SimilarityIndex
            A variable number of similarity indexes.
        """

        self.query_backends = query_backends

    @staticmethod
    def group_results_by_table(
        target_id: str,
        results: Iterable[Tuple[str, Iterable[float]]],
        table_groups: Optional[Dict] = None,
    ) -> Dict:
        """
        Groups column-based results by table.
        For a given query column, at most one candidate column is considered for each candidate table.
        This candidate column is the one with the highest sum of similarity scores.

        Parameters
        ----------
        target_id : str
            Typically the target column name used to get the results.
        results : Iterable[Tuple[str, Iterable[float]]]
            One or more pairs of column names (including the table names) and backend similarity scores.
        table_groups: Optional[Dict]
            Iteratively created table groups.
            If None, a new dict is created and populated with the current results.

        Returns
        -------
        Dict
            A mapping of table names to similarity scores.
        """

        if table_groups is None:
            table_groups = defaultdict(list)
        candidate_scores = {}
        for result_item, result_scores in results:
            name_components = result_item.split(".")
            table_name, column_name = (
                ".".join(name_components[:-1]),
                name_components[-1:][0],
            )

            candidate_column, existing_scores = candidate_scores.get(
                table_name, (None, None)
            )
            if existing_scores is None or sum(existing_scores) < sum(result_scores):
                candidate_scores[table_name] = (column_name, result_scores)

        for table_name, (candidate_column, result_scores) in candidate_scores.items():
            table_groups[table_name].append(
                ((target_id, candidate_column), result_scores)