text2sql-eval-toolkit 1.0.0__py3-none-any.whl

text2sql_eval_toolkit/logging.py

@@ -0,0 +1,54 @@
+#
+# Copyright IBM Corp. 2025 - 2026
+# SPDX-License-Identifier: Apache-2.0
+#
+
+import logging
+from tqdm import tqdm
+from pathlib import Path
+
+
+class TqdmLoggingHandler(logging.Handler):
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            tqdm.write(msg)
+            self.flush()
+        except Exception:
+            self.handleError(record)
+
+
+def get_logger(
+    name: str = "text2sql_eval_toolkit", level=logging.DEBUG, log_file: str = None
+):
+    logger = logging.getLogger(name)
+
+    if not logger.handlers:
+        logger.setLevel(level)
+
+        console_handler = TqdmLoggingHandler()
+        console_formatter = logging.Formatter(
+            "%(asctime)s - %(levelname)s - %(message)s", datefmt="%H:%M:%S"
+        )
+        console_handler.setFormatter(console_formatter)
+        logger.addHandler(console_handler)
+
+        # Default log file path relative to the project root
+        if log_file is None:
+            project_root = (
+                Path(__file__).resolve().parents[2]
+            )  # Go up from src/text2sql_eval_toolkit
+            log_file = project_root / "data" / "results" / "bak" / "log.txt"
+
+        log_file = Path(log_file)
+        log_file.parent.mkdir(parents=True, exist_ok=True)  # Ensure directory exists
+
+        # File handler
+        file_handler = logging.FileHandler(log_file, mode="w")
+        file_formatter = logging.Formatter(
+            "%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+        )
+        file_handler.setFormatter(file_formatter)
+        logger.addHandler(file_handler)
+
+    return logger

text2sql_eval_toolkit/profiling/profiling_tools.py

@@ -0,0 +1,185 @@
+#
+# Copyright IBM Corp. 2025 - 2026
+# SPDX-License-Identifier: Apache-2.0
+#
+
+import json
+import os
+from sqlglot import parse_one, exp
+import shutil
+from tqdm import tqdm
+from typing import Dict
+from text2sql_eval_toolkit.utils import get_gt_sqls
+from text2sql_eval_toolkit.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def analyze_sql_query(sql: str, dialect: str = "postgres") -> Dict:
+    """
+    Analyze a SQL query and classify it into categories with structural features and descriptive tags.
+
+    Args:
+        sql (str): The SQL query string.
+        dialect (str): SQL dialect for parsing (default is 'postgres').
+
+    Returns:
+        Dict: A dictionary with structural features and a set of descriptive tags.
+    """
+    parsed = parse_one(sql, dialect=dialect)
+
+    def count(exp_type):
+        return len(list(parsed.find_all(exp_type)))
+
+    def count_names(exp_type):
+        return len([e.name for e in parsed.find_all(exp_type)])
+
+    features = {
+        "query_table_count": count_names(exp.Table),
+        "query_column_count": count_names(exp.Column),
+        "query_nested_count": count(exp.Select)
+        + count(exp.Delete)
+        + count(exp.Insert)
+        - 1,
+        "query_aggregate_count": count(exp.AggFunc),
+        "query_sort_count": count(exp.Ordered),
+        "query_window_func_count": count(exp.Window),
+        "query_join_count": count(exp.Join),
+    }
+
+    # Classification logic
+    is_basic = (
+        features["query_table_count"] == 1
+        and features["query_nested_count"] == 0
+        and features["query_window_func_count"] == 0
+        and features["query_join_count"] == 0
+    )
+
+    is_multi_table = (
+        features["query_table_count"] > 1
+        and features["query_nested_count"] == 0
+        and features["query_window_func_count"] == 0
+        and features["query_join_count"] >= 1
+    )
+
+    is_advanced = features["query_table_count"] == 1 and (
+        features["query_nested_count"] > 0 or features["query_window_func_count"] > 0
+    )
+
+    # Tag generation
+    tags = set()
+    if is_basic:
+        tags.add("single_source_basic")
+    if is_multi_table:
+        tags.add("multi_table_simple")
+    if is_advanced:
+        tags.add("single_source_advanced")
+
+    if features["query_join_count"] > 0:
+        tags.add("has_join")
+    if features["query_nested_count"] > 0:
+        tags.add("has_nested_query")
+    if features["query_aggregate_count"] > 0:
+        tags.add("has_aggregation")
+    if features["query_sort_count"] > 0:
+        tags.add("has_sorting")
+    if features["query_window_func_count"] > 0:
+        tags.add("has_window_function")
+
+    return {"features": features, "categories": sorted(tags)}
+
+
+def merge_dictionaries(original_dict, new_dict):
+    """
+    Merges new_dict into original_dict in-place with specific logic for 'features' and 'categories' keys.
+
+    Args:
+        original_dict (dict): The original dictionary (modified in-place)
+        new_dict (dict): The new dictionary to merge
+
+    Returns:
+        bool: True if any conflicts/overwrites occurred, False otherwise
+    """
+    overwrite_occurred = False
+
+    for key, value in new_dict.items():
+        if key == "features":
+            # Initialize features if it doesn't exist
+            if key not in original_dict:
+                original_dict[key] = {}
+
+            # Merge features dictionaries
+            for feature_key, feature_value in value.items():
+                if (
+                    feature_key in original_dict[key]
+                    and original_dict[key][feature_key] != feature_value
+                ):
+                    overwrite_occurred = True
+                original_dict[key][feature_key] = feature_value
+
+        elif key == "categories":
+            # Initialize categories if it doesn't exist
+            if key not in original_dict:
+                original_dict[key] = []
+
+            # Add new categories that aren't already present
+            for category in value:
+                if category not in original_dict[key]:
+                    original_dict[key].append(category)
+
+        else:
+            # For other keys, new dictionary takes precedence
+            if key in original_dict and original_dict[key] != value:
+                overwrite_occurred = True
+            original_dict[key] = value
+
+    return overwrite_occurred
+
+
+def profile_pred_or_eval_json_file(
+    json_file_path: str, dialect: str = "postgres"
+) -> None:
+    # Load the JSON data
+    with open(json_file_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    backup_file_path = json_file_path + ".bak"
+    shutil.copy2(json_file_path, backup_file_path)
+
+    # Ensure the data is a list
+    if not isinstance(data, list):
+        raise ValueError("JSON file must contain an array of objects.")
+
+    # Track if any overwrites occur
+    overwrite_occurred = False
+    # Process each record in the input
+    for record in tqdm(data):
+        gt_sqls = get_gt_sqls(record)
+        sql_query = gt_sqls[0]
+        if len(gt_sqls) > 1:
+            logger.warning(
+                f"More than on gt query in record in {json_file_path}. Profiling only the first one."
+            )
+
+        try:
+            analysis_result = analyze_sql_query(sql_query, dialect)
+        except Exception as e:
+            logger.error(f"Failed to profile SQL query: {sql_query}. Error: {repr(e)}")
+            continue
+        # Initialize or update the 'meta' field
+        if "meta" not in record:
+            record["meta"] = analysis_result
+        else:
+            overwrite_occurred = merge_dictionaries(record["meta"], analysis_result)
+
+        # Write the updated data back to the original file
+        with open(json_file_path, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2)
+
+    if not overwrite_occurred:
+        os.remove(backup_file_path)
+    else:
+        print(f"Backup created at {backup_file_path} due to overwrites.")
+
+    logger.info(f"Profiling complete. Results written in {json_file_path}")

text2sql_eval_toolkit/utils.py

@@ -0,0 +1,302 @@
+#
+# Copyright IBM Corp. 2025 - 2026
+# SPDX-License-Identifier: Apache-2.0
+#
+
+import asyncio
+import importlib.resources as resources
+import json
+import os
+import time
+import pandas as pd
+from concurrent.futures import ThreadPoolExecutor, TimeoutError
+from typing import Any, Dict
+from pathlib import Path
+from text2sql_eval_toolkit.logging import get_logger
+
+
+BENCHMARKS_FILE = resources.files("text2sql_eval_toolkit.data").joinpath(
+    "benchmarks.json"
+)
+TEST_BENCHMARKS_FILE = resources.files("text2sql_eval_toolkit.data").joinpath(
+    "test-benchmarks.json"
+)
+logger = get_logger(__name__)
+
+
+def get_available_benchmarks(include_test: bool = True):
+    """
+    Get list of available benchmark IDs.
+    
+    Args:
+        include_test: If True, include test benchmarks from test-benchmarks.json
+    
+    Returns:
+        List of benchmark IDs
+    """
+    benchmarks = []
+    
+    # Load production benchmarks
+    if BENCHMARKS_FILE.exists():
+        with open(BENCHMARKS_FILE, "r") as f:
+            data = json.load(f)
+        benchmarks.extend(list(data.keys()))
+    
+    # Load test benchmarks if requested
+    if include_test and TEST_BENCHMARKS_FILE.exists():
+        with open(TEST_BENCHMARKS_FILE, "r") as f:
+            data = json.load(f)
+        benchmarks.extend(list(data.keys()))
+    
+    return benchmarks
+
+
+def get_benchmarks_info(is_test: bool = False) -> Dict[str, Any]:
+    """
+    Retrieves all the benchmarks' information.
+
+    Args:
+        is_test: If True, load test benchmarks from test-benchmarks.json
+
+    Returns:
+        Dict[str, Any]: Dictionary containing info and paths to benchmark files.
+    """
+    benchmarks_file = TEST_BENCHMARKS_FILE if is_test else BENCHMARKS_FILE
+    benchmarks_info = {}
+    try:
+        with open(benchmarks_file, "r") as meta_file:
+            benchmarks_meta = json.load(meta_file)
+    except Exception as e:
+        logger.error(f"Error loading the benchmarks JSON file: {benchmarks_file}.")
+        raise e
+    for benchmark_id, benchmark_info in benchmarks_meta.items():
+        root = BENCHMARKS_FILE.parent
+        if benchmark_id not in benchmarks_meta:
+            raise ValueError(
+                f"Benchmark ID '{benchmark_id}' not found in benchmarks.json."
+            )
+        benchmark_info = benchmarks_meta[benchmark_id]
+        benchmark_info["benchmark_json_path"] = resolve_path(
+            root, benchmark_info["data"]
+        )
+        benchmark_info["schema_json_path"] = resolve_path(
+            root, benchmark_info["schema"]
+        )
+        benchmark_info["predictions_path"] = resolve_path(
+            root, benchmark_info["predictions"]
+        )
+        benchmark_info["eval_results_path"] = Path(
+            benchmark_info["predictions_path"].with_name(
+                benchmark_info["predictions_path"].stem + "_eval.json"
+            )
+        ).resolve()
+        benchmark_info["eval_summary_path"] = Path(
+            benchmark_info["predictions_path"].with_name(
+                benchmark_info["predictions_path"].stem + "_eval_summary.json"
+            )
+        ).resolve()
+        benchmarks_info[benchmark_id] = benchmark_info
+    return benchmarks_info
+
+
+def resolve_path(root, path_str):
+    """
+    Resolves a given path string relative to a root directory.
+    If the provided path string is absolute, returns it as a Path object.
+    Otherwise, returns the path relative to the specified root directory.
+    Args:
+        root (Path): The root directory to resolve relative paths against.
+        path_str (str): The path string to resolve.
+    Returns:
+        Path: The resolved absolute or relative path as a Path object.
+    """
+    p = Path(path_str)
+    if p.is_absolute():
+        return p
+    return root / p
+
+
+def get_benchmark_info(benchmark_id: str, is_test: bool = False) -> Dict[str, Any]:
+    """
+    Retrieves the benchmark files for a given benchmark ID.
+    Automatically detects if benchmark is in test-benchmarks.json if not found in benchmarks.json.
+
+    Args:
+        benchmark_id (str): Identifier for the benchmark dataset.
+        is_test (bool): If True, load from test-benchmarks.json. If False, tries production first, then test.
+
+    Returns:
+        Dict[str, Any]: Dictionary containing info and paths to benchmark files.
+    """
+    # If is_test is True, only look in test benchmarks
+    if is_test:
+        benchmarks_file = TEST_BENCHMARKS_FILE
+        with open(benchmarks_file, "r") as meta_file:
+            benchmarks_meta = json.load(meta_file)
+        if benchmark_id not in benchmarks_meta:
+            raise ValueError(f"Benchmark ID '{benchmark_id}' not found in test-benchmarks.json.")
+    else:
+        # Try production benchmarks first
+        benchmarks_file = BENCHMARKS_FILE
+        with open(benchmarks_file, "r") as meta_file:
+            benchmarks_meta = json.load(meta_file)
+        
+        # If not found in production, try test benchmarks
+        if benchmark_id not in benchmarks_meta and TEST_BENCHMARKS_FILE.exists():
+            benchmarks_file = TEST_BENCHMARKS_FILE
+            with open(benchmarks_file, "r") as meta_file:
+                benchmarks_meta = json.load(meta_file)
+            if benchmark_id not in benchmarks_meta:
+                raise ValueError(f"Benchmark ID '{benchmark_id}' not found in benchmarks.json or test-benchmarks.json.")
+        elif benchmark_id not in benchmarks_meta:
+            raise ValueError(f"Benchmark ID '{benchmark_id}' not found in benchmarks.json.")
+    
+    root = benchmarks_file.parent
+    benchmark_info = benchmarks_meta[benchmark_id]
+    benchmark_info["benchmark_json_path"] = resolve_path(root, benchmark_info["data"])
+    benchmark_info["schema_json_path"] = resolve_path(root, benchmark_info["schema"])
+    benchmark_info["predictions_path"] = resolve_path(
+        root, benchmark_info["predictions"]
+    )
+    return benchmark_info
+
+
+def run_with_timeout(func, timeout=90, retries=2, wait=3, *args, **kwargs):
+    """
+    Runs a function with a timeout and retries.
+
+    Parameters:
+        func (callable): The function to run.
+        timeout (int): Timeout in seconds for each attempt.
+        retries (int): Number of retries after the first attempt.
+        wait (int): Seconds to wait between retries.
+        *args, **kwargs: Arguments to pass to the function.
+
+    Returns:
+        The result of the function if successful.
+
+    Raises:
+        TimeoutError: If all attempts time out.
+    """
+    for attempt in range(retries + 1):
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            future = executor.submit(func, *args, **kwargs)
+            try:
+                return future.result(timeout=timeout)
+            except TimeoutError:
+                logger.info(f"⚠️ Attempt {attempt + 1} timed out.")
+                if attempt < retries:
+                    time.sleep(wait)
+                else:
+                    raise TimeoutError(
+                        f"❗️ Function timed out after {retries + 1} attempts."
+                    )
+
+
+async def run_with_timeout_async(task, base_timeout=90, retries=2, wait=3):
+    for attempt in range(retries + 1):
+        timeout = base_timeout * (attempt + 1)
+        try:
+            return await asyncio.wait_for(task(), timeout=timeout)
+        except asyncio.TimeoutError:
+            logger.info(f"⚠️ Attempt {attempt + 1} timed out after {timeout} seconds.")
+            if attempt < retries:
+                logger.info(f"⏳ Retrying in {wait} seconds...")
+                await asyncio.sleep(wait)
+            else:
+                raise asyncio.TimeoutError(
+                    f"❗️ Function timed out after {retries + 1} attempts."
+                )
+
+
+def parse_dataframe(json_str):
+    """Reconstruct a DataFrame from a JSON-encoded dictionary."""
+    try:
+        df_dict = json.loads(json_str)
+        return pd.DataFrame(
+            data=df_dict["data"], columns=df_dict["columns"], index=df_dict["index"]
+        )
+    except Exception as e:
+        raise ValueError(
+            f"Failed to parse DataFrame JSON. Error: {e}. JSON string: {json_str}"
+        )
+
+
+def truncate_dataframe(
+    df: pd.DataFrame, head: int = 10, tail: int = 10
+) -> pd.DataFrame:
+    """
+    Truncate a DataFrame to show only the first `head` rows and last `tail` rows.
+    If the DataFrame has more rows than head + tail, insert a '...' row in between
+    with empty strings so print(df) looks clean.
+    """
+    if len(df) <= head + tail:
+        return df
+
+    top = df.head(head)
+    bottom = df.tail(tail)
+
+    # Row of empty strings with index labeled "..."
+    ellipsis_row = pd.DataFrame([[""] * df.shape[1]], columns=df.columns, index=["..."])
+
+    return pd.concat([top, ellipsis_row, bottom])
+
+
+def get_question_id(record):
+    """Gets Question ID from benchmark or prediction data record"""
+    id_keys = ["id", "question_id", "qid", "_id"]
+    for key in id_keys:
+        question_id = record.get(key)
+        if question_id is not None:
+            record["id"] = question_id  # Ensure 'id' is always set
+            return question_id
+    raise ValueError(f"Record has no ID field among {id_keys}: {record}")
+
+
+def get_utterance(record):
+    """Gets the question (utterance) from the benchmark or prediction data record"""
+    utterance_keys = ["utterance", "page_content", "question"]
+    for key in utterance_keys:
+        utterance = record.get(key)
+        if utterance:
+            record["utterance"] = utterance  # Ensure 'utterance' is always set
+            return utterance
+    raise ValueError(f"Record has no utterance field among {utterance_keys}: {record}")
+
+
+def get_gt_sqls(record):
+    gt_sql_keys = ["sql", "SQL", "target", "query"]
+    for key in gt_sql_keys:
+        gt_sqls = record.get(key)
+        if gt_sqls:
+            # Skip if it's a dict (structured SQL representation, not a string)
+            if isinstance(gt_sqls, dict):
+                continue
+            if not isinstance(gt_sqls, list):
+                gt_sqls = [gt_sqls]
+            record["sql"] = gt_sqls
+            return gt_sqls
+    if "metadata" in record and "sql" in record["metadata"]:
+        return [record["metadata"]["sql"]]
+    raise ValueError(f"Record has no ground truth SQL: {record}")
+
+
+def get_question(record):
+    return (
+        record["page_content"]
+        if "page_content" in record
+        else (record["question"] if "question" in record else record["utterance"])
+    )
+
+
+def get_default_eval_filename(predictions_file):
+    base_name, ext = os.path.splitext(predictions_file)
+    return f"{base_name}_eval{ext}"
+
+
+def add_summary_json_suffix(path: str) -> str:
+    return os.path.splitext(path)[0] + "_summary.json"
+
+
+def add_summary_csv_suffix(path: str) -> str:
+    return os.path.splitext(path)[0] + "_summary.csv"