collate-data-diff 0.11.2__py3-none-any.whl

data_diff/format.py

@@ -0,0 +1,359 @@
+import collections
+from enum import Enum
+from typing import Any, Optional, List, Dict, Tuple, Type
+
+import attrs
+from data_diff.diff_tables import DiffResultWrapper
+from data_diff.abcs.database_types import (
+    JSON,
+    Boolean,
+    ColType,
+    Array,
+    ColType_UUID,
+    Date,
+    FractionalType,
+    NumericType,
+    Struct,
+    TemporalType,
+    ColType_Alphanum,
+    String_Alphanum,
+)
+
+
+def jsonify_error(table1: List[str], table2: List[str], dbt_model: str, error: str) -> "FailedDiff":
+    return attrs.asdict(
+        FailedDiff(
+            status="failed",
+            model=dbt_model,
+            dataset1=table1,
+            dataset2=table2,
+            error=error,
+        )
+    )
+
+
+Columns = List[Tuple[str, str, ColType]]
+
+
+def jsonify(
+    diff: DiffResultWrapper,
+    dbt_model: str,
+    dataset1_columns: Columns,
+    dataset2_columns: Columns,
+    columns_diff: Dict[str, List[str]],
+    with_summary: bool = False,
+    stats_only: bool = False,
+) -> "JsonDiff":
+    """
+    Converts the diff result into a JSON-serializable format.
+    Optionally add stats summary and schema diff.
+    """
+    diff_info = diff.info_tree.info
+    table1 = diff_info.tables[0]
+    table2 = diff_info.tables[1]
+    key_columns = table1.key_columns
+
+    t1_exclusive_rows = []
+    t2_exclusive_rows = []
+    diff_rows = []
+    rows = None
+    schema = [field for field, _ in diff_info.diff_schema]
+
+    t1_exclusive_rows, t2_exclusive_rows, diff_rows = _group_rows(diff_info, schema)
+
+    if not stats_only:
+        rows = _make_rows_diff(t1_exclusive_rows, t2_exclusive_rows, diff_rows, key_columns)
+
+    summary = None
+    if with_summary:
+        summary = _jsonify_diff_summary(diff.get_stats_dict(is_dbt=True))
+
+    columns = _jsonify_columns_diff(dataset1_columns, dataset2_columns, columns_diff, list(key_columns))
+
+    is_different = bool(
+        t1_exclusive_rows
+        or t2_exclusive_rows
+        or diff_rows
+        or (columns_diff["added"] or columns_diff["removed"] or columns_diff["changed"])
+    )
+    return attrs.asdict(
+        JsonDiff(
+            status="success",
+            result="different" if is_different else "identical",
+            model=dbt_model,
+            dataset1=list(table1.table_path),
+            dataset2=list(table2.table_path),
+            rows=rows,
+            summary=summary,
+            columns=columns,
+        )
+    )
+
+
+@attrs.define(frozen=True)
+class JsonExclusiveRowValue:
+    """
+    Value of a single column in a row
+    """
+
+    isPK: bool
+    value: Any
+
+
+@attrs.define(frozen=True)
+class JsonDiffRowValue:
+    """
+    Pair of diffed values for 2 rows with equal PKs
+    """
+
+    dataset1: Any
+    dataset2: Any
+    isDiff: bool
+    isPK: bool
+
+
+@attrs.define(frozen=True)
+class Total:
+    dataset1: int
+    dataset2: int
+
+
+@attrs.define(frozen=True)
+class ExclusiveRows:
+    dataset1: int
+    dataset2: int
+
+
+@attrs.define(frozen=True)
+class Rows:
+    total: Total
+    exclusive: ExclusiveRows
+    updated: int
+    unchanged: int
+
+
+@attrs.define(frozen=True)
+class Stats:
+    diffCounts: Dict[str, int]
+
+
+@attrs.define(frozen=True)
+class JsonDiffSummary:
+    rows: Rows
+    stats: Stats
+
+
+@attrs.define(frozen=True)
+class ExclusiveColumns:
+    dataset1: List[str]
+    dataset2: List[str]
+
+
+class ColumnKind(Enum):
+    INTEGER = "integer"
+    FLOAT = "float"
+    STRING = "string"
+    DATE = "date"
+    TIME = "time"
+    DATETIME = "datetime"
+    BOOL = "boolean"
+    UNSUPPORTED = "unsupported"
+
+
+KIND_MAPPING: List[Tuple[Type[ColType], ColumnKind]] = [
+    (Boolean, ColumnKind.BOOL),
+    (Date, ColumnKind.DATE),
+    (TemporalType, ColumnKind.DATETIME),
+    (FractionalType, ColumnKind.FLOAT),
+    (NumericType, ColumnKind.INTEGER),
+    (ColType_UUID, ColumnKind.STRING),
+    (ColType_Alphanum, ColumnKind.STRING),
+    (String_Alphanum, ColumnKind.STRING),
+    (JSON, ColumnKind.STRING),
+    (Array, ColumnKind.STRING),
+    (Struct, ColumnKind.STRING),
+    (ColType, ColumnKind.UNSUPPORTED),
+]
+
+
+@attrs.define(frozen=True)
+class Column:
+    name: str
+    type: str
+    kind: str
+
+
+@attrs.define(frozen=True)
+class JsonColumnsSummary:
+    dataset1: List[Column]
+    dataset2: List[Column]
+    primaryKey: List[str]
+    exclusive: ExclusiveColumns
+    typeChanged: List[str]
+
+
+@attrs.define(frozen=True)
+class ExclusiveDiff:
+    dataset1: List[Dict[str, JsonExclusiveRowValue]]
+    dataset2: List[Dict[str, JsonExclusiveRowValue]]
+
+
+@attrs.define(frozen=True)
+class RowsDiff:
+    exclusive: ExclusiveDiff
+    diff: List[Dict[str, JsonDiffRowValue]]
+
+
+@attrs.define(frozen=True)
+class FailedDiff:
+    status: str  # Literal ["failed"]
+    model: str
+    dataset1: List[str]
+    dataset2: List[str]
+    error: str
+
+    version: str = "1.0.0"
+
+
+@attrs.define(frozen=True)
+class JsonDiff:
+    status: str  # Literal ["success"]
+    result: str  # Literal ["different", "identical"]
+    model: str
+    dataset1: List[str]
+    dataset2: List[str]
+    rows: Optional[RowsDiff]
+    summary: Optional[JsonDiffSummary]
+    columns: Optional[JsonColumnsSummary]
+
+    version: str = "1.1.0"
+
+
+def _group_rows(
+    diff_info: DiffResultWrapper, schema: List[str]
+) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]:
+    t1_exclusive_rows = []
+    t2_exclusive_rows = []
+    diff_rows = []
+
+    for row in diff_info.diff:
+        row_w_schema = dict(zip(schema, row))
+        is_t1_exclusive = row_w_schema["is_exclusive_a"]
+        is_t2_exclusive = row_w_schema["is_exclusive_b"]
+
+        if is_t1_exclusive:
+            t1_exclusive_rows.append(row_w_schema)
+
+        elif is_t2_exclusive:
+            t2_exclusive_rows.append(row_w_schema)
+
+        else:
+            diff_rows.append(row_w_schema)
+
+    return t1_exclusive_rows, t2_exclusive_rows, diff_rows
+
+
+def _make_rows_diff(
+    t1_exclusive_rows: List[Dict[str, Any]],
+    t2_exclusive_rows: List[Dict[str, Any]],
+    diff_rows: List[Dict[str, Any]],
+    key_columns: List[str],
+) -> RowsDiff:
+    diff_rows_jsonified = []
+    for row in diff_rows:
+        diff_rows_jsonified.append(_jsonify_diff(row, key_columns))
+
+    t1_exclusive_rows_jsonified = []
+    for row in t1_exclusive_rows:
+        t1_exclusive_rows_jsonified.append(_jsonify_exclusive(row, key_columns))
+
+    t2_exclusive_rows_jsonified = []
+    for row in t2_exclusive_rows:
+        t2_exclusive_rows_jsonified.append(_jsonify_exclusive(row, key_columns))
+
+    return RowsDiff(
+        exclusive=ExclusiveDiff(dataset1=t1_exclusive_rows_jsonified, dataset2=t2_exclusive_rows_jsonified),
+        diff=diff_rows_jsonified,
+    )
+
+
+def _jsonify_diff(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, JsonDiffRowValue]:
+    columns = collections.defaultdict(dict)
+    for field, value in row.items():
+        if field in ("is_exclusive_a", "is_exclusive_b"):
+            continue
+
+        if field.startswith("is_diff_"):
+            column_name = field[len("is_diff_") :]
+            columns[column_name]["isDiff"] = bool(value)
+
+        elif field.endswith("_a"):
+            column_name = field[: -len("_a")]
+            columns[column_name]["dataset1"] = value
+            columns[column_name]["isPK"] = column_name in key_columns
+
+        elif field.endswith("_b"):
+            column_name = field[: -len("_b")]
+            columns[column_name]["dataset2"] = value
+            columns[column_name]["isPK"] = column_name in key_columns
+
+    return {column: JsonDiffRowValue(**data) for column, data in columns.items()}
+
+
+def _jsonify_exclusive(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, JsonExclusiveRowValue]:
+    columns = collections.defaultdict(dict)
+    for field, value in row.items():
+        if field in ("is_exclusive_a", "is_exclusive_b"):
+            continue
+        if field.startswith("is_diff_"):
+            continue
+        if field.endswith("_b") and row["is_exclusive_b"]:
+            column_name = field[: -len("_b")]
+            columns[column_name]["isPK"] = column_name in key_columns
+            columns[column_name]["value"] = value
+        elif field.endswith("_a") and row["is_exclusive_a"]:
+            column_name = field[: -len("_a")]
+            columns[column_name]["isPK"] = column_name in key_columns
+            columns[column_name]["value"] = value
+    return {column: JsonExclusiveRowValue(**data) for column, data in columns.items()}
+
+
+def _jsonify_diff_summary(stats_dict: dict) -> JsonDiffSummary:
+    return JsonDiffSummary(
+        rows=Rows(
+            total=Total(dataset1=stats_dict["rows_A"], dataset2=stats_dict["rows_B"]),
+            exclusive=ExclusiveRows(
+                dataset1=stats_dict["exclusive_A"],
+                dataset2=stats_dict["exclusive_B"],
+            ),
+            updated=stats_dict["updated"],
+            unchanged=stats_dict["unchanged"],
+        ),
+        stats=Stats(diffCounts=stats_dict["values"]),
+    )
+
+
+def _jsonify_columns_diff(
+    dataset1_columns: Columns, dataset2_columns: Columns, columns_diff: Dict[str, List[str]], key_columns: List[str]
+) -> JsonColumnsSummary:
+    return JsonColumnsSummary(
+        dataset1=[
+            Column(name=name, type=type_, kind=_map_kind(kind).value) for (name, type_, kind) in dataset1_columns
+        ],
+        dataset2=[
+            Column(name=name, type=type_, kind=_map_kind(kind).value) for (name, type_, kind) in dataset2_columns
+        ],
+        primaryKey=key_columns,
+        exclusive=ExclusiveColumns(
+            dataset2=list(columns_diff.get("added", [])),
+            dataset1=list(columns_diff.get("removed", [])),
+        ),
+        typeChanged=list(columns_diff.get("changed", [])),
+    )
+
+
+def _map_kind(kind: ColType) -> ColumnKind:
+    for raw_kind, json_kind in KIND_MAPPING:
+        if isinstance(kind, raw_kind):
+            return json_kind
+    return ColumnKind.UNSUPPORTED

data_diff/hashdiff_tables.py

@@ -0,0 +1,264 @@
+import os
+from numbers import Number
+import logging
+from collections import defaultdict
+from typing import Any, Collection, Dict, Iterator, List, Sequence, Set, Tuple
+
+import attrs
+from typing_extensions import Literal
+
+from data_diff.abcs.database_types import ColType_UUID, NumericType, PrecisionType, StringType, Boolean, JSON
+from data_diff.info_tree import InfoTree
+from data_diff.utils import safezip, diffs_are_equiv_jsons
+from data_diff.thread_utils import ThreadedYielder
+from data_diff.table_segment import TableSegment
+from data_diff.diff_tables import TableDiffer
+
+BENCHMARK = os.environ.get("BENCHMARK", False)
+
+DEFAULT_BISECTION_THRESHOLD = 1024 * 16
+DEFAULT_BISECTION_FACTOR = 32
+
+logger = logging.getLogger("hashdiff_tables")
+
+# Just for local readability: TODO: later switch to real type declarations of these.
+_Op = Literal["+", "-"]
+_PK = Sequence[Any]
+_Row = Tuple[Any]
+
+
+def diff_sets(
+    a: Sequence[_Row],
+    b: Sequence[_Row],
+    *,
+    json_cols: dict = None,
+    columns1: Sequence[str],
+    columns2: Sequence[str],
+    key_columns1: Sequence[str],
+    key_columns2: Sequence[str],
+    ignored_columns1: Collection[str],
+    ignored_columns2: Collection[str],
+) -> Iterator:
+    # Group full rows by PKs on each side. The first items are the PK: TableSegment.relevant_columns
+    rows_by_pks1: Dict[_PK, List[_Row]] = defaultdict(list)
+    rows_by_pks2: Dict[_PK, List[_Row]] = defaultdict(list)
+    for row in a:
+        pk: _PK = tuple(val for col, val in zip(key_columns1, row))
+        rows_by_pks1[pk].append(row)
+    for row in b:
+        pk: _PK = tuple(val for col, val in zip(key_columns2, row))
+        rows_by_pks2[pk].append(row)
+
+    # Mind that the same pk MUST go in full with all the -/+ rows all at once, for grouping.
+    diffs_by_pks: Dict[_PK, List[Tuple[_Op, _Row]]] = defaultdict(list)
+    for pk in sorted(set(rows_by_pks1) | set(rows_by_pks2)):
+        cutrows1: List[_Row] = [
+            tuple(val for col, val in zip(columns1, row1) if col not in ignored_columns1) for row1 in rows_by_pks1[pk]
+        ]
+        cutrows2: List[_Row] = [
+            tuple(val for col, val in zip(columns2, row2) if col not in ignored_columns2) for row2 in rows_by_pks2[pk]
+        ]
+
+        # Either side has 0 rows: a clearly exclusive row.
+        # Either side has 2+ rows: duplicates on either side, yield it all regardless of values.
+        # Both sides == 1: non-duplicate, non-exclusive, so check for values of interest.
+        if len(cutrows1) != 1 or len(cutrows2) != 1 or cutrows1 != cutrows2:
+            for row1 in rows_by_pks1[pk]:
+                diffs_by_pks[pk].append(("-", row1))
+            for row2 in rows_by_pks2[pk]:
+                diffs_by_pks[pk].append(("+", row2))
+
+    warned_diff_cols = set()
+    for diffs in (diffs_by_pks[pk] for pk in sorted(diffs_by_pks)):
+        if json_cols:
+            parsed_match, overriden_diff_cols = diffs_are_equiv_jsons(diffs, json_cols)
+            if parsed_match:
+                to_warn = overriden_diff_cols - warned_diff_cols
+                for w in to_warn:
+                    logger.warning(
+                        f"Equivalent JSON objects with different string representations detected "
+                        f"in column '{w}'. These cases are NOT reported as differences."
+                    )
+                    warned_diff_cols.add(w)
+                continue
+        yield from diffs
+
+
+@attrs.define(frozen=False)
+class HashDiffer(TableDiffer):
+    """Finds the diff between two SQL tables
+
+    The algorithm uses hashing to quickly check if the tables are different, and then applies a
+    bisection search recursively to find the differences efficiently.
+
+    Works best for comparing tables that are mostly the same, with minor discrepancies.
+
+    Parameters:
+        bisection_factor (int): Into how many segments to bisect per iteration.
+        bisection_threshold (Number): When should we stop bisecting and compare locally (in row count).
+        threaded (bool): Enable/disable threaded diffing. Needed to take advantage of database threads.
+        max_threadpool_size (int): Maximum size of each threadpool. ``None`` means auto.
+                                   Only relevant when `threaded` is ``True``.
+                                   There may be many pools, so number of actual threads can be a lot higher.
+    """
+
+    bisection_factor: int = DEFAULT_BISECTION_FACTOR
+    bisection_threshold: int = DEFAULT_BISECTION_THRESHOLD
+    bisection_disabled: bool = False  # i.e. always download the rows (used in tests)
+
+    stats: dict = attrs.field(factory=dict)
+
+    def __attrs_post_init__(self) -> None:
+        # Validate options
+        if self.bisection_factor >= self.bisection_threshold:
+            raise ValueError("Incorrect param values (bisection factor must be lower than threshold)")
+        if self.bisection_factor < 2:
+            raise ValueError("Must have at least two segments per iteration (i.e. bisection_factor >= 2)")
+
+    def _validate_and_adjust_columns(self, table1: TableSegment, table2: TableSegment, *, strict: bool = True) -> None:
+        for c1, c2 in safezip(table1.relevant_columns, table2.relevant_columns):
+            if c1 not in table1._schema:
+                raise ValueError(f"Column '{c1}' not found in schema for table {table1}")
+            if c2 not in table2._schema:
+                raise ValueError(f"Column '{c2}' not found in schema for table {table2}")
+
+            # Update schemas to minimal mutual precision
+            col1 = table1._schema[c1]
+            col2 = table2._schema[c2]
+            if isinstance(col1, PrecisionType):
+                if not isinstance(col2, PrecisionType):
+                    if strict:
+                        raise TypeError(f"Incompatible types for column '{c1}':  {col1} <-> {col2}")
+                    else:
+                        continue
+
+                lowest = min(col1, col2, key=lambda col: col.precision)
+
+                if col1.precision != col2.precision:
+                    logger.warning(f"Using reduced precision {lowest} for column '{c1}'. Types={col1}, {col2}")
+
+                table1._schema[c1] = attrs.evolve(col1, precision=lowest.precision, rounds=lowest.rounds)
+                table2._schema[c2] = attrs.evolve(col2, precision=lowest.precision, rounds=lowest.rounds)
+
+            elif isinstance(col1, (NumericType, Boolean)):
+                if not isinstance(col2, (NumericType, Boolean)):
+                    if strict:
+                        raise TypeError(f"Incompatible types for column '{c1}':  {col1} <-> {col2}")
+                    else:
+                        continue
+
+                lowest = min(col1, col2, key=lambda col: col.precision)
+
+                if col1.precision != col2.precision:
+                    logger.warning(f"Using reduced precision {lowest} for column '{c1}'. Types={col1}, {col2}")
+
+                if lowest.precision != col1.precision:
+                    table1._schema[c1] = attrs.evolve(col1, precision=lowest.precision)
+                if lowest.precision != col2.precision:
+                    table2._schema[c2] = attrs.evolve(col2, precision=lowest.precision)
+
+        for t in [table1, table2]:
+            for c in t.relevant_columns:
+                ctype = t._schema[c]
+                if not ctype.supported:
+                    logger.warning(
+                        f"[{t.database.name}] Column '{c}' of type '{ctype}' has no compatibility handling. "
+                        "If encoding/formatting differs between databases, it may result in false positives."
+                    )
+
+    def _diff_segments(
+        self,
+        ti: ThreadedYielder,
+        table1: TableSegment,
+        table2: TableSegment,
+        info_tree: InfoTree,
+        max_rows: int,
+        level=0,
+        segment_index=None,
+        segment_count=None,
+    ):
+        logger.info(
+            ". " * level + f"Diffing segment {segment_index}/{segment_count}, "
+            f"key-range: {table1.min_key}..{table2.max_key}, "
+            f"size <= {max_rows}"
+        )
+
+        # When benchmarking, we want the ability to skip checksumming. This
+        # allows us to download all rows for comparison in performance. By
+        # default, data-diff will checksum the section first (when it's below
+        # the threshold) and _then_ download it.
+        if BENCHMARK:
+            if self.bisection_disabled or max_rows < self.bisection_threshold:
+                return self._bisect_and_diff_segments(ti, table1, table2, info_tree, level=level, max_rows=max_rows)
+
+        (count1, checksum1), (count2, checksum2) = self._threaded_call("count_and_checksum", [table1, table2])
+
+        assert not info_tree.info.rowcounts
+        info_tree.info.rowcounts = {1: count1, 2: count2}
+
+        if count1 == 0 and count2 == 0:
+            logger.debug(
+                "Uneven distribution of keys detected in segment %s..%s (big gaps in the key column). "
+                "For better performance, we recommend to increase the bisection-threshold.",
+                table1.min_key,
+                table1.max_key,
+            )
+            assert checksum1 is None and checksum2 is None
+            info_tree.info.is_diff = False
+            return
+
+        if checksum1 == checksum2:
+            info_tree.info.is_diff = False
+            return
+
+        info_tree.info.is_diff = True
+        return self._bisect_and_diff_segments(ti, table1, table2, info_tree, level=level, max_rows=max(count1, count2))
+
+    def _bisect_and_diff_segments(
+        self,
+        ti: ThreadedYielder,
+        table1: TableSegment,
+        table2: TableSegment,
+        info_tree: InfoTree,
+        level=0,
+        max_rows=None,
+    ):
+        assert table1.is_bounded and table2.is_bounded
+
+        max_space_size = max(table1.approximate_size(), table2.approximate_size())
+        if max_rows is None:
+            # We can be sure that row_count <= max_rows iff the table key is unique
+            max_rows = max_space_size
+            info_tree.info.max_rows = max_rows
+
+        # If count is below the threshold, just download and compare the columns locally
+        # This saves time, as bisection speed is limited by ping and query performance.
+        if self.bisection_disabled or max_rows < self.bisection_threshold or max_space_size < self.bisection_factor * 2:
+            rows1, rows2 = self._threaded_call("get_values", [table1, table2])
+            json_cols = {
+                i: colname
+                for i, colname in enumerate(table1.extra_columns)
+                if isinstance(table1._schema[colname], JSON)
+            }
+            diff = list(
+                diff_sets(
+                    rows1,
+                    rows2,
+                    json_cols=json_cols,
+                    columns1=table1.relevant_columns,
+                    columns2=table2.relevant_columns,
+                    key_columns1=table1.key_columns,
+                    key_columns2=table2.key_columns,
+                    ignored_columns1=self.ignored_columns1,
+                    ignored_columns2=self.ignored_columns2,
+                )
+            )
+
+            info_tree.info.set_diff(diff)
+            info_tree.info.rowcounts = {1: len(rows1), 2: len(rows2)}
+
+            logger.info(". " * level + f"Diff found {len(diff)} different rows.")
+            self.stats["rows_downloaded"] = self.stats.get("rows_downloaded", 0) + max(len(rows1), len(rows2))
+            return diff
+
+        return super()._bisect_and_diff_segments(ti, table1, table2, info_tree, level, max_rows)

data_diff/info_tree.py

@@ -0,0 +1,62 @@
+from typing import List, Dict, Optional, Any, Tuple, Union
+
+import attrs
+from typing_extensions import Self
+
+from data_diff.table_segment import TableSegment
+
+
+@attrs.define(frozen=False)
+class SegmentInfo:
+    tables: List[TableSegment]
+
+    diff: Optional[List[Union[Tuple[Any, ...], List[Any]]]] = None
+    diff_schema: Optional[Tuple[Tuple[str, type], ...]] = None
+    is_diff: Optional[bool] = None
+    diff_count: Optional[int] = None
+
+    rowcounts: Dict[int, int] = attrs.field(factory=dict)
+    max_rows: Optional[int] = None
+
+    def set_diff(
+        self, diff: List[Union[Tuple[Any, ...], List[Any]]], schema: Optional[Tuple[Tuple[str, type]]] = None
+    ) -> None:
+        self.diff_schema = schema
+        self.diff = diff
+        self.diff_count = len(diff)
+        self.is_diff = self.diff_count > 0
+
+    def update_from_children(self, child_infos) -> None:
+        child_infos = list(child_infos)
+        assert child_infos
+
+        # self.diff = list(chain(*[c.diff for c in child_infos]))
+        self.diff_count = sum(c.diff_count for c in child_infos if c.diff_count is not None)
+        self.is_diff = any(c.is_diff for c in child_infos)
+        self.diff_schema = next((child.diff_schema for child in child_infos if child.diff_schema is not None), None)
+        self.diff = sum((c.diff for c in child_infos if c.diff is not None), [])
+
+        self.rowcounts = {
+            1: sum(c.rowcounts[1] for c in child_infos if c.rowcounts),
+            2: sum(c.rowcounts[2] for c in child_infos if c.rowcounts),
+        }
+
+
+@attrs.define(frozen=True)
+class InfoTree:
+    SEGMENT_INFO_CLASS = SegmentInfo
+
+    info: SegmentInfo
+    children: List["InfoTree"] = attrs.field(factory=list)
+
+    def add_node(self, table1: TableSegment, table2: TableSegment, max_rows: Optional[int] = None) -> Self:
+        cls = self.__class__
+        node = cls(cls.SEGMENT_INFO_CLASS([table1, table2], max_rows=max_rows))
+        self.children.append(node)
+        return node
+
+    def aggregate_info(self) -> None:
+        if self.children:
+            for c in self.children:
+                c.aggregate_info()
+            self.info.update_from_children(c.info for c in self.children)