clickzetta-semantic-model-generator 1.0.1__tar.gz → 1.0.2__tar.gz

{clickzetta_semantic_model_generator-1.0.1 → clickzetta_semantic_model_generator-1.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: clickzetta-semantic-model-generator
-Version: 1.0.1
+Version: 1.0.2
 Summary: Curate a Semantic Model for ClickZetta Lakehouse
 License: Apache Software License; BSD License
 Author: qililiang

{clickzetta_semantic_model_generator-1.0.1 → clickzetta_semantic_model_generator-1.0.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "clickzetta-semantic-model-generator"
-version = "1.0.1"
+version = "1.0.2"
 description = "Curate a Semantic Model for ClickZetta Lakehouse"
 authors = ["qililiang <qililiang@clickzetta.com>"]
 license = "Apache Software License; BSD License"

clickzetta_semantic_model_generator-1.0.2/semantic_model_generator/relationships/__init__.py

@@ -0,0 +1,15 @@
+"""Public APIs for relationship discovery."""
+
+from .discovery import (
+    RelationshipDiscoveryResult,
+    RelationshipSummary,
+    discover_relationships_from_schema,
+    discover_relationships_from_tables,
+)
+
+__all__ = [
+    "RelationshipDiscoveryResult",
+    "RelationshipSummary",
+    "discover_relationships_from_schema",
+    "discover_relationships_from_tables",
+]

clickzetta_semantic_model_generator-1.0.2/semantic_model_generator/relationships/discovery.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any, Iterable, List, Optional, Sequence, Tuple
+
+import pandas as pd
+from loguru import logger
+
+from semantic_model_generator.clickzetta_utils.clickzetta_connector import (
+    _TABLE_NAME_COL,
+    _TABLE_SCHEMA_COL,
+    get_table_representation,
+    get_valid_schemas_tables_columns_df,
+)
+from semantic_model_generator.data_processing import data_types
+from semantic_model_generator.data_processing.data_types import FQNParts, Table
+from semantic_model_generator.generate_model import (
+    _DEFAULT_N_SAMPLE_VALUES_PER_COL,
+    _infer_relationships,
+)
+from semantic_model_generator.protos import semantic_model_pb2
+
+try:  # pragma: no cover - optional dependency for type checking
+    from clickzetta.zettapark.session import Session
+except Exception:  # pragma: no cover
+    Session = Any  # type: ignore
+
+DEFAULT_MAX_WORKERS = 4
+
+
+@dataclass
+class RelationshipSummary:
+    total_tables: int
+    total_columns: int
+    total_relationships_found: int
+    processing_time_ms: int
+
+
+@dataclass
+class RelationshipDiscoveryResult:
+    relationships: List[semantic_model_pb2.Relationship]
+    tables: List[Table]
+    summary: RelationshipSummary
+
+
+def _normalize_table_names(table_names: Optional[Iterable[str]]) -> Optional[List[str]]:
+    if table_names is None:
+        return None
+    return [name.upper() for name in table_names]
+
+
+def _build_tables_from_dataframe(
+    session: Session,
+    workspace: str,
+    schema: str,
+    columns_df: pd.DataFrame,
+    sample_values_per_column: int,
+    max_workers: int = DEFAULT_MAX_WORKERS,
+) -> List[Tuple[FQNParts, Table]]:
+    if columns_df.empty:
+        return []
+
+    if _TABLE_NAME_COL not in columns_df.columns:
+        raise KeyError(
+            f"Expected '{_TABLE_NAME_COL}' column in metadata dataframe. "
+            "Ensure information_schema query returned table names."
+        )
+
+    table_order = (
+        columns_df[_TABLE_NAME_COL]
+        .astype(str)
+        .str.upper()
+        .drop_duplicates()
+        .tolist()
+    )
+
+    tables: List[Tuple[FQNParts, Table]] = []
+    for idx, table_name in enumerate(table_order):
+        table_columns_df = columns_df[columns_df[_TABLE_NAME_COL] == table_name]
+        if table_columns_df.empty:
+            continue
+
+        max_workers_for_table = min(max_workers, len(table_columns_df.index) or 1)
+        table_proto = get_table_representation(
+            session=session,
+            workspace=workspace,
+            schema_name=schema,
+            table_name=table_name,
+            table_index=idx,
+            ndv_per_column=sample_values_per_column,
+            columns_df=table_columns_df,
+            max_workers=max_workers_for_table,
+        )
+        tables.append(
+            (
+                FQNParts(database=workspace, schema_name=schema, table=table_name),
+                table_proto,
+            )
+        )
+
+    return tables
+
+
+def _discover_relationships(
+    raw_tables: List[Tuple[FQNParts, Table]],
+    strict_join_inference: bool,
+    session: Optional[Session],
+) -> List[semantic_model_pb2.Relationship]:
+    if not raw_tables:
+        return []
+
+    relationships = _infer_relationships(
+        raw_tables,
+        session=session if strict_join_inference else None,
+        strict_join_inference=strict_join_inference,
+    )
+    return relationships
+
+
+def discover_relationships_from_tables(
+    tables: Sequence[Tuple[FQNParts, Table]],
+    *,
+    strict_join_inference: bool = False,
+    session: Optional[Session] = None,
+) -> RelationshipDiscoveryResult:
+    """
+    Run relationship inference using pre-constructed table metadata.
+    """
+    start = time.perf_counter()
+    relationships = _discover_relationships(
+        list(tables),
+        strict_join_inference=strict_join_inference,
+        session=session,
+    )
+    end = time.perf_counter()
+
+    all_columns = sum(len(table.columns) for _, table in tables)
+    summary = RelationshipSummary(
+        total_tables=len(tables),
+        total_columns=all_columns,
+        total_relationships_found=len(relationships),
+        processing_time_ms=int((end - start) * 1000),
+    )
+
+    return RelationshipDiscoveryResult(
+        relationships=relationships,
+        tables=[table for _, table in tables],
+        summary=summary,
+    )
+
+
+def discover_relationships_from_schema(
+    session: Session,
+    workspace: str,
+    schema: str,
+    *,
+    table_names: Optional[Sequence[str]] = None,
+    sample_values_per_column: int = _DEFAULT_N_SAMPLE_VALUES_PER_COL,
+    strict_join_inference: bool = False,
+    max_workers: int = DEFAULT_MAX_WORKERS,
+) -> RelationshipDiscoveryResult:
+    """
+    Discover table relationships for all tables in a ClickZetta schema.
+    """
+    normalized_tables = _normalize_table_names(table_names)
+
+    metadata_df = get_valid_schemas_tables_columns_df(
+        session=session,
+        workspace=workspace,
+        table_schema=schema,
+        table_names=normalized_tables,
+    )
+    metadata_df.columns = [str(col).upper() for col in metadata_df.columns]
+
+    if metadata_df.empty:
+        logger.warning(
+            "No column metadata found for workspace=%s schema=%s tables=%s",
+            workspace,
+            schema,
+            table_names,
+        )
+        return RelationshipDiscoveryResult(
+            relationships=[],
+            tables=[],
+            summary=RelationshipSummary(
+                total_tables=0,
+                total_columns=0,
+                total_relationships_found=0,
+                processing_time_ms=0,
+            ),
+        )
+
+    raw_tables = _build_tables_from_dataframe(
+        session=session,
+        workspace=workspace,
+        schema=schema,
+        columns_df=metadata_df,
+        sample_values_per_column=sample_values_per_column,
+        max_workers=max_workers,
+    )
+
+    return discover_relationships_from_tables(
+        raw_tables,
+        strict_join_inference=strict_join_inference,
+        session=session,
+    )

clickzetta_semantic_model_generator-1.0.2/semantic_model_generator/tests/relationship_discovery_test.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+import pandas as pd
+
+from semantic_model_generator.relationships.discovery import (
+    discover_relationships_from_schema,
+)
+
+
+class _FakeResult:
+    def __init__(self, df: pd.DataFrame):
+        self._df = df
+
+    def to_pandas(self) -> pd.DataFrame:
+        return self._df.copy()
+
+
+class _FakeSession:
+    def __init__(self, tables: List[str], columns_df: pd.DataFrame):
+        self.tables = tables
+        self.columns_df = columns_df
+
+    def sql(self, query: str):
+        normalized = query.upper()
+        if "SHOW CATALOGS" in normalized:
+            return _FakeResult(
+                pd.DataFrame(
+                    {
+                        "CATALOG_NAME": ["CLICKZETTA_SAMPLE_DATA"],
+                        "CATEGORY": ["MANAGED"],
+                    }
+                )
+            )
+        if "FROM INFORMATION_SCHEMA.TABLES" in normalized:
+            data = {"TABLE_SCHEMA": ["TPCH_100G"] * len(self.tables), "TABLE_NAME": self.tables}
+            return _FakeResult(pd.DataFrame(data))
+        if "FROM INFORMATION_SCHEMA.COLUMNS" in normalized:
+            return _FakeResult(self.columns_df)
+        if "SELECT DISTINCT" in normalized:
+            # Return single column of sample values
+            return _FakeResult(pd.DataFrame({"VALUE": [1, 2, 3]}))
+        raise AssertionError(f"Unexpected query: {query}")
+
+
+def _build_columns_df() -> pd.DataFrame:
+    records: List[Dict[str, Any]] = []
+    # Orders table
+    records.extend(
+        [
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "ORDERS",
+                "COLUMN_NAME": "ORDER_ID",
+                "DATA_TYPE": "NUMBER",
+                "IS_PRIMARY_KEY": True,
+            },
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "ORDERS",
+                "COLUMN_NAME": "CUSTOMER_ID",
+                "DATA_TYPE": "NUMBER",
+                "IS_PRIMARY_KEY": False,
+            },
+        ]
+    )
+    # Customer table
+    records.extend(
+        [
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "CUSTOMER",
+                "COLUMN_NAME": "CUSTOMER_ID",
+                "DATA_TYPE": "NUMBER",
+                "IS_PRIMARY_KEY": True,
+            },
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "CUSTOMER",
+                "COLUMN_NAME": "NAME",
+                "DATA_TYPE": "STRING",
+                "IS_PRIMARY_KEY": False,
+            },
+        ]
+    )
+    return pd.DataFrame.from_records(records)
+
+
+def test_discover_relationships_from_schema_builds_relationships():
+    tables = ["ORDERS", "CUSTOMER"]
+    columns_df = _build_columns_df()
+    session = _FakeSession(tables, columns_df)
+
+    result = discover_relationships_from_schema(
+        session=session,
+        workspace="CLICKZETTA_SAMPLE_DATA",
+        schema="TPCH_100G",
+        strict_join_inference=False,
+    )
+
+    assert result.summary.total_tables == 2
+    assert result.summary.total_relationships_found >= 1
+
+    names = {rel.name for rel in result.relationships}
+    assert any("ORDERS" in name and "CUSTOMER" in name for name in names)
+
+    left_tables = {rel.left_table for rel in result.relationships}
+    right_tables = {rel.right_table for rel in result.relationships}
+    assert "ORDERS" in left_tables
+    assert "CUSTOMER" in right_tables