clickzetta-semantic-model-generator 1.0.1__py3-none-any.whl → 1.0.2__py3-none-any.whl

{clickzetta_semantic_model_generator-1.0.1.dist-info → clickzetta_semantic_model_generator-1.0.2.dist-info}/METADATA

@@ -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.dist-info → clickzetta_semantic_model_generator-1.0.2.dist-info}/RECORD

@@ -16,10 +16,13 @@ semantic_model_generator/output_models/.keep,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRk
 semantic_model_generator/protos/semantic_model.proto,sha256=WZiN4b8vR-ZX-Lj9Vsm6HjZNAyNvM1znIyut_YkPVSI,16473
 semantic_model_generator/protos/semantic_model_pb2.py,sha256=scbWkW-I-r3_hp_5SHoOWn02p52RJ9DJ0_-nRgr0LHc,25606
 semantic_model_generator/protos/semantic_model_pb2.pyi,sha256=iiBIZxtX9d6IuUO3aLcsJsHUeZqdi14vYNuUsSM8C0g,18267
+semantic_model_generator/relationships/__init__.py,sha256=HN6Opie25Oawt2fCDM_bZwRBVBEzqRsEXgDzYC7ytns,373
+semantic_model_generator/relationships/discovery.py,sha256=fy8mzfAZtWWVUjiSP5jeKffoa2GodNIKA3eGi4dxhHo,6020
 semantic_model_generator/tests/clickzetta_connector_test.py,sha256=NKDbhll8TxFtbwFJIUcQQDKJU1LkfEz9nL0SE32hx3o,3114
 semantic_model_generator/tests/cte_utils_test.py,sha256=8v2nrrD2GkH_PTGIsKm3lQ06unzis8iR31atT8bUX98,17385
 semantic_model_generator/tests/generate_model_classification_test.py,sha256=q7dh29h9iF17ChzuoFSLtRX9ASiAm2oY4OkGyFVfn5Y,2117
 semantic_model_generator/tests/llm_enrichment_test.py,sha256=YeYg4voQ3wy2vgF7H9JNdnMOyZmUfiMdL6oEXFv-ztg,14415
+semantic_model_generator/tests/relationship_discovery_test.py,sha256=1SVX59-mpHQvxk7RDGtglesg6VXU9TnnZZcfjZi5IHs,3448
 semantic_model_generator/tests/relationships_filters_test.py,sha256=fVyA-hwxGqdlFr_PGT9YdrPz13XGvI4J_5F3IpVTFEE,8009
 semantic_model_generator/tests/samples/validate_yamls.py,sha256=262j-2i2oFZtTyK2susOrbxxE5eS-6IN-V0jFEOpt_w,156249
 semantic_model_generator/tests/utils_test.py,sha256=Tfvb-ErZPBS_HjXr4N7XSJnL3hlncNGF5pay1xFfNHg,539
@@ -29,7 +32,7 @@ semantic_model_generator/validate/context_length.py,sha256=HL-GfaRXNcVji1-pAFGXG
 semantic_model_generator/validate/keywords.py,sha256=mFtmIK72YLQ7wODL-zCC-uMXzXMJT6Tg4fISAin2WtQ,5480
 semantic_model_generator/validate/schema.py,sha256=ImFUzO5EzNbl65-0dqt5AmdEvy4lqX8gBkOSOKN8uZM,5863
 semantic_model_generator/validate_model.py,sha256=Uq-V-GfPeF2Dy4l9uF5Guv104gDCDGh0Cxz1AJOu5dk,836
-clickzetta_semantic_model_generator-1.0.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-clickzetta_semantic_model_generator-1.0.1.dist-info/METADATA,sha256=mbWjvg1PtkZlQfPAE06kHDe6whzheZ0iYj_HfXNqqGE,7816
-clickzetta_semantic_model_generator-1.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-clickzetta_semantic_model_generator-1.0.1.dist-info/RECORD,,
+clickzetta_semantic_model_generator-1.0.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+clickzetta_semantic_model_generator-1.0.2.dist-info/METADATA,sha256=vgObGkoBnyfxUuJM8eAH8ZnIGuJrueJFnU1UfTo_gN4,7816
+clickzetta_semantic_model_generator-1.0.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+clickzetta_semantic_model_generator-1.0.2.dist-info/RECORD,,

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",
+]

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,
+    )

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