arthur-common 1.0.1__py3-none-any.whl

arthur_common/tools/aggregation_analyzer.py

@@ -0,0 +1,243 @@
+import inspect
+import logging
+import typing
+import uuid
+from typing import Any, Callable, get_type_hints
+
+from arthur_common.aggregations import (
+    AggregationFunction,
+    NumericAggregationFunction,
+    SketchAggregationFunction,
+)
+from arthur_common.models.metrics import (
+    AggregationMetricType,
+    AggregationSpecSchema,
+    DatasetReference,
+    MetricsColumnParameterSchema,
+    MetricsDatasetParameterSchema,
+    MetricsLiteralParameterSchema,
+    MetricsParameterSchemaUnion,
+)
+from arthur_common.models.schema_definitions import (
+    DType,
+    MetricColumnParameterAnnotation,
+    MetricDatasetParameterAnnotation,
+    MetricLiteralParameterAnnotation,
+    MetricsParameterAnnotationUnion,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FunctionAnalyzer:
+    @staticmethod
+    def _python_type_to_scope_dtype(t: Any) -> DType:
+        if t is int:
+            return DType.INT
+        elif t is str:
+            return DType.STRING
+        elif t is bool:
+            return DType.BOOL
+        elif t is float:
+            return DType.FLOAT
+        elif t is uuid.UUID:
+            return DType.UUID
+        elif t is DatasetReference:
+            return DType.UUID
+        else:
+            raise ValueError(f"Parameter type {t} is not supported.")
+
+    @staticmethod
+    def _get_metric_annotation_from_annotated(
+        param_name: str,
+        annotation: typing.Annotated,  # type: ignore
+    ) -> MetricsParameterAnnotationUnion | None:
+        arthur_metric_annotations = [
+            m
+            for m in annotation.__metadata__
+            if isinstance(m, MetricsParameterAnnotationUnion)
+        ]
+        metric_annotation: MetricsParameterAnnotationUnion | None = None
+        if len(arthur_metric_annotations) == 1:
+            metric_annotation = arthur_metric_annotations[0]
+        if len(arthur_metric_annotations) > 1:
+            raise ValueError(
+                f"Parameter {param_name} defines more than one metric annotation.",
+            )
+        return metric_annotation
+
+    @staticmethod
+    def _get_scope_metric_parameter_from_annotation(
+        param_name: str,
+        param_dtype: DType,
+        optional: bool,
+        annotation: typing.Annotated,  # type: ignore
+    ) -> MetricsParameterSchemaUnion:
+        if annotation is None:
+            return MetricsLiteralParameterSchema(
+                parameter_key=param_name,
+                optional=optional,
+                parameter_dtype=param_dtype,
+                friendly_name=param_name,
+                description=f"A {param_dtype.value} value.",
+            )
+        elif isinstance(annotation, MetricLiteralParameterAnnotation):
+            return MetricsLiteralParameterSchema(
+                parameter_key=param_name,
+                optional=optional,
+                parameter_dtype=param_dtype,
+                friendly_name=annotation.friendly_name,
+                description=annotation.description,
+            )
+        elif isinstance(annotation, MetricDatasetParameterAnnotation):
+            if param_dtype != DType.UUID:
+                raise ValueError(
+                    f"Dataset parameter {param_name} has type {param_dtype}, but should be a UUID.",
+                )
+            return MetricsDatasetParameterSchema(
+                parameter_key=param_name,
+                optional=optional,
+                friendly_name=annotation.friendly_name,
+                description=annotation.description,
+                model_problem_type=annotation.model_problem_type,
+            )
+        elif isinstance(annotation, MetricColumnParameterAnnotation):
+            if param_dtype != DType.STRING:
+                raise ValueError(
+                    f"Column parameter {param_name} has type {param_dtype}, but should be a string.",
+                )
+            return MetricsColumnParameterSchema(
+                parameter_key=param_name,
+                tag_hints=annotation.tag_hints,
+                optional=optional,
+                source_dataset_parameter_key=annotation.source_dataset_parameter_key,
+                allowed_column_types=annotation.allowed_column_types,
+                allow_any_column_type=annotation.allow_any_column_type,
+                friendly_name=annotation.friendly_name,
+                description=annotation.description,
+            )
+        else:
+            raise ValueError(
+                f"Parameter {param_name} has an unsupported annotation {annotation}.",
+            )
+
+    """
+    Returns a list of parameter names, parameter types, scope-specific annotations.
+    """
+
+    @staticmethod
+    def _extract_parameter_metadata(func: Callable) -> list[MetricsParameterSchemaUnion]:  # type: ignore
+        parameter_schemas: list[MetricsParameterSchemaUnion] = []
+        args = inspect.signature(func).parameters
+        for name, param in args.items():
+            if name == "self":
+                continue
+            if name == "ddb_conn":
+                continue
+
+            if param.annotation == inspect.Parameter.empty:
+                raise ValueError(
+                    f"{func.__name__} must provide type annotation for parameter {name}.",
+                )
+            parameter_schemas.append(
+                FunctionAnalyzer._get_scope_metric_parameter_from_annotation(
+                    name,
+                    FunctionAnalyzer._python_type_to_scope_dtype(
+                        get_type_hints(func)[name],
+                    ),
+                    param.default != inspect.Parameter.empty,
+                    (
+                        FunctionAnalyzer._get_metric_annotation_from_annotated(
+                            name,
+                            param.annotation,
+                        )
+                        if typing.get_origin(param.annotation) is typing.Annotated
+                        else None
+                    ),
+                ),
+            )
+
+        return parameter_schemas
+
+    @staticmethod
+    def analyze_aggregation_function(agg_func: type) -> AggregationSpecSchema:
+        # Check if X is a subclass of AggregationFunction
+        if not issubclass(agg_func, AggregationFunction):
+            raise TypeError(
+                f"Class {agg_func.__name__} is not a subclass of AggregationFunction.",
+            )
+
+        if issubclass(agg_func, NumericAggregationFunction):
+            metric_type = AggregationMetricType.NUMERIC
+        elif issubclass(agg_func, SketchAggregationFunction):
+            metric_type = AggregationMetricType.SKETCH
+        else:
+            raise ValueError(
+                f"Class {agg_func.__name__} is not a subclass of SketchAggregationFunction, NumericAggregationFunction.",
+            )
+        # Check if X implements the required methods
+        required_methods = ["aggregate", "id", "description", "display_name"]
+        static_methods = ["description", "id", "display_name"]
+        for method in required_methods:
+            if not hasattr(agg_func, method) or not callable(getattr(agg_func, method)):
+                raise AttributeError(
+                    f"Class {agg_func.__name__} does not implement {method} method.",
+                )
+
+        for method in static_methods:
+            if not is_static_method(getattr(agg_func, method)):
+                raise AttributeError(f"Method {method} should be a staticmethod.")
+        # Check if X passes the ABC implementation:
+        try:
+            agg_func()
+        except TypeError as e:
+            if "Can't instantiate abstract class" in str(e):
+                logger.error(str(e))
+                raise TypeError(
+                    f"Class {agg_func.__name__} does not implement all the base class functions.",
+                )
+            else:
+                # This is okay, it just means we didn't supply proper args to the __init__ function. The ABC mismatch would throw before this, so it must have passed
+                pass
+
+        aggregation_init_args: list[MetricsParameterSchemaUnion] = []
+
+        # This is necessary because all the way down in the ABC class, some __init__ function is defined which we don't care about. Users should be able to exclude an init function and this allows them to do that.
+        if has_custom_init(agg_func):
+            aggregation_init_args = FunctionAnalyzer._extract_parameter_metadata(
+                agg_func.__init__,
+            )
+        aggregate_args = FunctionAnalyzer._extract_parameter_metadata(
+            agg_func.aggregate,
+        )
+
+        aggregation_id = agg_func.id()
+        aggregation_description = agg_func.description()
+
+        return AggregationSpecSchema(
+            name=agg_func.display_name(),
+            id=aggregation_id,
+            # TODO: Require description, version
+            description=aggregation_description,
+            # version=0,
+            metric_type=metric_type,
+            init_args=aggregation_init_args,
+            aggregate_args=aggregate_args,
+        )
+
+
+def has_custom_init(cls: type) -> bool:
+    init_method = getattr(cls, "__init__", None)
+    base_init_method = (
+        getattr(cls.__base__, "__init__", None) if hasattr(cls, "__base__") else None
+    )
+    return init_method is not base_init_method
+
+
+def is_static_method(method: type) -> bool:
+    if inspect.isfunction(method):
+        # Check if the method accepts no arguments or only default arguments
+        argspec = inspect.getfullargspec(method)
+        if len(argspec.args) == 0 and not argspec.varargs and not argspec.varkw:
+            return True
+    return False

arthur_common/tools/aggregation_loader.py

@@ -0,0 +1,59 @@
+import inspect
+import logging
+from types import ModuleType
+from typing import Type
+
+import arthur_common.aggregations as agg_module
+from arthur_common.aggregations.aggregator import (
+    AggregationFunction,
+    NumericAggregationFunction,
+    SketchAggregationFunction,
+)
+from arthur_common.models.metrics import AggregationSpecSchema
+from arthur_common.tools.aggregation_analyzer import FunctionAnalyzer
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class AggregationLoader:
+    @staticmethod
+    def load_aggregations() -> (
+        list[tuple[AggregationSpecSchema, Type[AggregationFunction]]]
+    ):
+        def find_subclasses(
+            module: ModuleType,
+            base_classes: tuple[type, ...],
+            visited: set[ModuleType] = set(),
+        ) -> set[type]:
+            subclasses = set()
+            visited.add(module)
+            for name, obj in inspect.getmembers(module):
+                if inspect.isclass(obj) and issubclass(obj, base_classes):
+                    subclasses.add(obj)
+                elif inspect.ismodule(obj) and obj not in visited:
+                    subclasses.update(find_subclasses(obj, base_classes, visited))
+            return subclasses
+
+        base_classes = (SketchAggregationFunction, NumericAggregationFunction)
+        agg_functions = find_subclasses(agg_module, base_classes)
+        aggregation_specs = []
+        """
+          This seems to pick up duplicate functions somehow in different namespaces, ie:
+           <class 'categorical_count.CategoricalCountAggregationFunction'>
+           and
+           <class 'arthur_common.aggregations.functions.categorical_count.CategoricalCountAggregationFunction'>
+           so dedupe by id
+        """
+        aggregation_ids = set()
+        for agg_function in agg_functions:
+            try:
+                func_spec = FunctionAnalyzer.analyze_aggregation_function(agg_function)
+                logger.info(f"Found agg function {agg_function}")
+                if func_spec.id in aggregation_ids:
+                    continue
+                aggregation_specs.append((func_spec, agg_function))
+                aggregation_ids.add(func_spec.id)
+            except Exception as e:
+                logger.error(f"Failed to load aggregation function {agg_function}: {e}")
+        return aggregation_specs

arthur_common/tools/duckdb_data_loader.py

@@ -0,0 +1,329 @@
+import json
+from typing import Any
+
+import duckdb
+import pandas as pd
+from arthur_common.models.datasets import DatasetJoinKind
+from arthur_common.models.schema_definitions import (
+    DatasetListType,
+    DatasetObjectType,
+    DatasetScalarType,
+    DatasetSchema,
+    DType,
+)
+from dateutil.parser import parse
+from fsspec import filesystem
+from pydantic import BaseModel
+
+
+class ColumnFormat(BaseModel):
+    source_name: str
+    alias: str
+    format: str
+
+
+class DuckDBOperator:
+    """
+    Loads data into a DuckDB table.
+
+    If no schema is supplied, the output table will contain columns with names equal to the source names in the data.
+    If a schema is applied, the column names in the output table will be aliases equal to the column id from the schema.
+    This allows for consistent column naming across different data sources.
+    """
+
+    @staticmethod
+    def load_data_to_duckdb(
+        data: list[dict[str, Any]] | pd.DataFrame,
+        preprocess_schema: bool = False,
+        table_name: str = "inferences",
+        conn: duckdb.DuckDBPyConnection | None = None,
+        schema: DatasetSchema | None = None,
+    ) -> duckdb.DuckDBPyConnection:
+        if not conn:
+            conn = duckdb.connect()
+
+        if type(data) == list:
+            DuckDBOperator._load_unstructured_data(data, table_name, conn, schema)
+        elif type(data) == pd.DataFrame:
+            DuckDBOperator._load_structured_data(
+                data,
+                preprocess_schema,
+                table_name,
+                conn,
+                schema,
+            )
+        else:
+            raise ValueError(f"Unsupported data type: {type(data)}")
+        return conn
+
+    """
+    Rename columns from ids to their friendly names based on schema.column_names
+    """
+
+    @staticmethod
+    def apply_alias_mask(
+        table_name: str,
+        conn: duckdb.DuckDBPyConnection,
+        schema: DatasetSchema,
+    ) -> None:
+        old_new_mask = {
+            str(col_id): schema.column_names[col_id] for col_id in schema.column_names
+        }
+        DuckDBOperator._apply_alias_mask(table_name, conn, old_new_mask)
+
+    @staticmethod
+    def _apply_alias_mask(
+        table_name: str,
+        conn: duckdb.DuckDBPyConnection,
+        old_new_mask: dict[str, str],
+    ) -> None:
+        for old, new in old_new_mask.items():
+            # Don't quote the join column names, since they're already quoted as part of escape_identifier's output
+            alter_query = f"ALTER TABLE {table_name} RENAME COLUMN {escape_identifier(old)} TO {escape_identifier(new)}"
+            conn.sql(alter_query)
+
+    @staticmethod
+    def _load_unstructured_data(
+        data: list[dict[str, Any]],
+        table_name: str,
+        conn: duckdb.DuckDBPyConnection,
+        schema: DatasetSchema | None,
+    ) -> None:
+        with filesystem("memory").open(f"inferences.json", "w") as file:
+            file.write(json.dumps(data))
+        conn.register_filesystem(filesystem("memory"))
+
+        if schema:
+            column_formats = make_duckdb_dataset_schema(schema)
+
+            key_value_pairs = [
+                f"{escape_identifier(col.source_name)}: '{col.format}'"
+                for col in column_formats
+            ]
+            stringified_schema = ", ".join([f"{kv}" for kv in key_value_pairs])
+            stringified_schema = f"{{ {stringified_schema} }}"
+
+            read_stmt = f"read_json('memory://inferences.json', format='array', columns={stringified_schema})"
+        else:
+            read_stmt = "read_json_auto('memory://inferences.json')"
+
+        conn.sql(
+            f"CREATE OR REPLACE TEMP TABLE {table_name} AS SELECT * FROM {read_stmt}",
+        )
+
+        if schema:
+            old_new_mask = {}
+            for col in column_formats:
+                old_new_mask[col.source_name] = col.alias
+            DuckDBOperator._apply_alias_mask(table_name, conn, old_new_mask)
+
+    @staticmethod
+    def _load_structured_data(
+        data: pd.DataFrame,
+        preprocess_schema: bool,
+        table_name: str,
+        conn: duckdb.DuckDBPyConnection,
+        schema: DatasetSchema | None,
+    ) -> None:
+        if preprocess_schema:
+            data = DuckDBOperator._preprocess_dataframe_schema_inference(data)
+
+        if schema:
+            column_formats = make_duckdb_dataset_schema(schema)
+
+            key_value_pairs = [
+                f"{escape_identifier(col.source_name)} {col.format}"
+                for col in column_formats
+            ]
+            stringified_schema = ", ".join([f"{kv}" for kv in key_value_pairs])
+            create_table_stmt = (
+                f"CREATE OR REPLACE TEMP TABLE {table_name} ({stringified_schema});"
+            )
+            conn.sql(create_table_stmt)
+            conn.sql(f"INSERT INTO {table_name} SELECT * FROM data")
+
+            old_new_mask = {}
+            for col in column_formats:
+                old_new_mask[col.source_name] = col.alias
+            DuckDBOperator._apply_alias_mask(table_name, conn, old_new_mask)
+
+        else:
+            conn.sql(f"CREATE OR REPLACE TEMP TABLE {table_name} AS SELECT * FROM data")
+
+    """
+        Preprocess to make smarter type inferences. Pandas and json recognize very little beyond primitives out of the box. We can support a little more with a little effort like:
+        1. Datetimes
+
+        Modifies the input data in place to have smarter types than pandas will natively infer
+    """
+
+    @staticmethod
+    def _preprocess_dataframe_schema_inference(data: pd.DataFrame) -> pd.DataFrame:
+        datetime_columns = _infer_dataframe_datetime_columns(data)
+        for column in datetime_columns:
+            try:
+                data[column] = pd.to_datetime(data[column])
+            except Exception:
+                # we're using best-effort to infer datetime columns, but just in case we got it wrong, move on
+                continue
+
+        return data
+
+    @staticmethod
+    def join_tables(
+        conn: duckdb.DuckDBPyConnection,
+        table_name: str,
+        table_1: str,
+        table_2: str,
+        table_1_join_key: str,
+        table_2_join_key: str,
+        join_kind: DatasetJoinKind = DatasetJoinKind.INNER,
+    ) -> None:
+        match join_kind:
+            case DatasetJoinKind.INNER:
+                join = "INNER"
+            case DatasetJoinKind.LEFT_OUTER:
+                join = "LEFT"
+            case DatasetJoinKind.RIGHT_OUTER:
+                join = "RIGHT"
+            case DatasetJoinKind.OUTER:
+                join = "FULL OUTER"
+            case _:
+                raise NotImplementedError(f"Join kind {join_kind} is not supported.")
+
+        # Don't quote the join column names, since they're already quoted as part of escape_identifier's output
+        join_query = f"""
+        CREATE TABLE {table_name} AS
+        SELECT *
+        FROM {table_1} a
+        {join} JOIN {table_2} b
+        ON a.{escape_identifier(table_1_join_key)} = b.{escape_identifier(table_2_join_key)}
+        """
+
+        conn.sql(join_query)
+
+
+def _infer_dataframe_datetime_columns(df: pd.DataFrame, n: int = 100) -> list[str]:
+    """
+    Infer datetime columns in a pandas DataFrame by parsing non-null values in the first n rows. Return the column names believed to be datetime type
+
+    Parameters:
+        df (pandas.DataFrame): Input DataFrame.
+        n (int): Number of non-null rows to consider for each column. Default is 100.
+
+    Returns:
+        datetime_columns (list): List of column names inferred to be datetime.
+    """
+    datetime_columns = []
+
+    for column in df.columns:
+        non_null_values = df[column].dropna().head(n)
+        if non_null_values.empty:
+            continue
+
+        # Try parsing each non-null value in the column
+        try:
+            parsed_values = non_null_values.apply(lambda x: parse(x))
+
+            # If parsing succeeds for all values, consider the column as datetime
+            if parsed_values.notnull().all():
+                datetime_columns.append(column)
+        except:
+            # If parsing fails for any value, move to the next column
+            continue
+
+    return datetime_columns
+
+
+"""
+Returns a list of ColumnFormat. Depending on structure / unstructured data, we need to format the root columns differently, so return the raw forms.
+
+See the subtle differences between
+
+CREATE TABLE users (
+    userID BIGINT,
+    userName VARCHAR,
+    hobbies ARRAY<VARCHAR>
+);
+
+and
+
+SELECT *
+FROM read_json('todos.json',
+               format = 'array',
+               columns = {userId: 'UBIGINT',
+                          userName: 'VARCHAR',
+                          hobbies: 'ARRAY<VARCHAR>'});
+
+"""
+
+
+def make_duckdb_dataset_schema(schema: DatasetSchema) -> list[ColumnFormat]:
+    details = []
+    for col in schema.columns:
+        format = _make_schema(col.definition)
+        details.append(
+            ColumnFormat(source_name=col.source_name, alias=str(col.id), format=format),
+        )
+
+    return details
+
+
+def _make_schema(
+    schema_node: DatasetObjectType | DatasetListType | DatasetScalarType,
+) -> str:
+    if isinstance(schema_node, DatasetObjectType):
+        details = {}
+        for col, value in schema_node.object.items():
+            details[col] = _make_schema(value)
+        key_value_pairs = [
+            f"{escape_identifier(col)} {value}" for col, value in details.items()
+        ]
+        return f"STRUCT({', '.join(key_value_pairs)})"
+
+    elif isinstance(schema_node, DatasetListType):
+        return f"{_make_schema(schema_node.items)}[]"
+    elif isinstance(schema_node, DatasetScalarType):
+        match schema_node.dtype:
+            case DType.INT:
+                return "BIGINT"
+            case DType.FLOAT:
+                return "DOUBLE"
+            case DType.BOOL:
+                return "BOOLEAN"
+            case DType.STRING:
+                return "VARCHAR"
+            case DType.UUID:
+                return "UUID"
+            case DType.TIMESTAMP:
+                return "TIMESTAMP"
+            case DType.JSON:
+                return "JSON"
+            case _:
+                raise ValueError(f"Unknown mapping for DType {schema_node.dtype}")
+    else:
+        raise NotImplementedError(
+            f"Schema conversion not implemented for node type {type(schema_node)}",
+        )
+
+
+def escape_identifier(identifier: str) -> str:
+    """
+    Escape an identifier (e.g., column name) for use in a SQL query.
+    This method handles special characters and ensures proper quoting.
+    """
+    # Replace any double quotes with two double quotes
+    escaped = identifier.replace('"', '""')
+    # Wrap the entire identifier in double quotes and return
+    return f'"{escaped}"'
+
+
+def escape_str_literal(literal: str) -> str:
+    """
+    Escape a duckDB string literal for use in a SQL query.
+    https://duckdb.org/docs/stable/sql/data_types/literal_types.html#escape-string-literals
+    """
+    # replace any single quotes with two single quotes
+    escaped = literal.replace("'", "''")
+    # Wrap the entire identifier in single quotes and return
+    return f"'{escaped}'"

arthur_common/tools/functions.py

@@ -0,0 +1,46 @@
+import hashlib
+from uuid import UUID
+
+"""
+Convert a uuid to a 12 character, all lowercase string deterministically
+"""
+
+
+def uuid_to_base26(uuid: str | UUID) -> str:
+    # Remove hyphens from the UUID
+    if isinstance(uuid, UUID):
+        uuid_str = str(uuid)
+    else:
+        uuid_str = uuid
+
+    no_hyphens = uuid_str.replace("-", "")
+
+    # Convert the hex string to an integer
+    num = int(no_hyphens, 16)
+
+    # Define the alphabet for base-26
+    alphabet = "abcdefghijklmnopqrstuvwxyz"
+    base = len(alphabet)
+
+    # Encode the integer into a base-26 string
+    base26_str = ""
+    while num > 0:
+        num, rem = divmod(num, base)
+        base26_str = alphabet[rem] + base26_str
+
+    # Ensure the string is 12 characters long (pad with 'a' if necessary)
+    base26_str = base26_str.rjust(12, "a")
+
+    # If the resulting string is longer than 12 characters, take the last 12 characters
+    return base26_str[-12:]
+
+
+"""
+Hash a string
+"""
+
+
+def hash_nonce(nonce: str) -> str:
+    md5_hash = hashlib.md5()
+    md5_hash.update(nonce.encode("utf-8"))
+    return md5_hash.hexdigest()