data-designer 0.1.0__py3-none-any.whl

data_designer/engine/processing/gsonschema/validators.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from copy import deepcopy
+import logging
+from typing import Any, overload
+
+from jsonschema import Draft202012Validator, ValidationError, validators
+
+from data_designer.engine.processing.gsonschema.exceptions import JSONSchemaValidationError
+from data_designer.engine.processing.gsonschema.schema_transformers import forbid_additional_properties
+from data_designer.engine.processing.gsonschema.types import DataObjectT, JSONSchemaT, T_primitive
+
+DEFAULT_JSONSCHEMA_VALIDATOR = Draft202012Validator
+
+logger = logging.getLogger(__name__)
+
+
+def prune_additional_properties(
+    _, allow_additional_properties: bool, instance: DataObjectT, schema: JSONSchemaT
+) -> None:
+    """A JSONSchemaValidtor extension function to prune additional properties.
+
+    Operates on an individual schema in-place.
+
+    Args:
+        allow_additional_properties (bool): The value of the `additionalProperties`
+            field for this schema.
+        instance (DataObjectT): The data object being validated.
+        schema (JSONSchemaT): The schema for this object.
+
+    Returns:
+        Nothing (in place)
+    """
+    # Only act if the instance is a dict.
+    if not isinstance(instance, dict) or allow_additional_properties:
+        return
+
+    # Allowed keys are those defined in the schema's "properties".
+    allowed = schema.get("properties", {}).keys()
+
+    # Iterate over a copy of keys so we can modify the dict in place.
+    n_removed = 0
+    for key in list(instance.keys()):
+        if key not in allowed:
+            instance.pop(key)
+            n_removed += 1
+            logger.info(f"Unspecified property removed from data object: {key}.")
+
+    if n_removed > 0:
+        logger.info(f"{n_removed} unspecified properties removed from data object.")
+
+
+def extend_jsonschema_validator_with_pruning(validator):
+    """Modify behavior of a jsonschema.Validator to use pruning.
+
+    Validators extended using this function will prune extra
+    fields, rather than raising a ValidationError, when encountering
+    extra, unspecified fiends when `additionalProperties: False` is
+    set in the validating schema.
+
+    Args:
+        validator (Type[jsonschema.Validator): A validator class
+            to extend with pruning behavior.
+
+    Returns:
+        Type[jsonschema.Validator]: A validator class that will
+            prune extra fields.
+    """
+    return validators.extend(validator, {"additionalProperties": prune_additional_properties})
+
+
+## We don't expect the outer data type (e.g. dict, list, or const) to be
+## modified by the pruning action.
+@overload
+def validate(
+    obj: dict[str, Any],
+    schema: JSONSchemaT,
+    pruning: bool = False,
+    no_extra_properties: bool = False,
+) -> dict[str, Any]: ...
+
+
+@overload
+def validate(
+    obj: list[Any],
+    schema: JSONSchemaT,
+    pruning: bool = False,
+    no_extra_properties: bool = False,
+) -> list[Any]: ...
+
+
+@overload
+def validate(
+    obj: T_primitive,
+    schema: JSONSchemaT,
+    pruning: bool = False,
+    no_extra_properties: bool = False,
+) -> T_primitive: ...
+
+
+def validate(
+    obj: DataObjectT,
+    schema: JSONSchemaT,
+    pruning: bool = False,
+    no_extra_properties: bool = False,
+) -> DataObjectT:
+    """Validate a data object against a JSONSchema.
+
+    Args:
+        obj (DataObjectT): A data structure to validate against the
+            schema.
+        schema: (JSONSchemaT): A valid JSONSchema to use to validate
+            the provided data object.
+        pruning (bool): If set to `True`, then the default behavior for
+            `additionalProperties: False` is set to prune non-specified
+            properties instead of raising a ValidationError.
+            Default: `False`.
+        no_extra_properties (bool): If set to `True`, then
+            `additionalProperties: False` is set on all the schema
+            and all of its sub-schemas. This operation overrides any
+            existing settings of `additionalProperties` within the
+            schema. Default: `False`.
+
+    Raises:
+        JSONSchemaValidationError: This exception raised in the
+            event that the JSONSchema doesn't match the provided
+            schema.
+    """
+    final_object = deepcopy(obj)
+    validator = DEFAULT_JSONSCHEMA_VALIDATOR
+    if pruning:
+        validator = extend_jsonschema_validator_with_pruning(validator)
+
+    if no_extra_properties:
+        schema = forbid_additional_properties(schema)
+
+    try:
+        validator(schema).validate(final_object)
+    except ValidationError as exc:
+        raise JSONSchemaValidationError(str(exc)) from exc
+
+    return final_object

data_designer/engine/processing/processors/base.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, DataT, TaskConfigT
+
+
+class Processor(ConfigurableTask[TaskConfigT], ABC):
+    @staticmethod
+    @abstractmethod
+    def metadata() -> ConfigurableTaskMetadata: ...
+
+    @abstractmethod
+    def process(self, data: DataT, *, current_batch_number: int | None = None) -> DataT: ...

data_designer/engine/processing/processors/drop_columns.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import pandas as pd
+
+from data_designer.config.processors import DropColumnsProcessorConfig
+from data_designer.engine.configurable_task import ConfigurableTaskMetadata
+from data_designer.engine.dataset_builders.artifact_storage import BatchStage
+from data_designer.engine.processing.processors.base import Processor
+
+logger = logging.getLogger(__name__)
+
+
+class DropColumnsProcessor(Processor[DropColumnsProcessorConfig]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="drop_columns",
+            description="Drop columns from the input dataset.",
+            required_resources=None,
+        )
+
+    def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None) -> pd.DataFrame:
+        logger.info(f"🙈 Dropping columns: {self.config.column_names}")
+        if current_batch_number is not None:  # not in preview mode
+            self._save_dropped_columns_if_needed(data, current_batch_number)
+        for column in self.config.column_names:
+            if column in data.columns:
+                data.drop(columns=[column], inplace=True)
+            else:
+                logger.warning(f"⚠️ Cannot drop column: `{column}` not found in the dataset.")
+        return data
+
+    def _save_dropped_columns_if_needed(self, data: pd.DataFrame, current_batch_number: int) -> None:
+        logger.debug("📦 Saving dropped columns to dropped-columns directory")
+        dropped_column_parquet_file_name = self.artifact_storage.create_batch_file_path(
+            batch_number=current_batch_number,
+            batch_stage=BatchStage.DROPPED_COLUMNS,
+        ).name
+        self.artifact_storage.write_parquet_file(
+            parquet_file_name=dropped_column_parquet_file_name,
+            dataframe=data[self.config.column_names],
+            batch_stage=BatchStage.DROPPED_COLUMNS,
+        )

data_designer/engine/processing/processors/registry.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.processors import (
+    DropColumnsProcessorConfig,
+    ProcessorType,
+)
+from data_designer.engine.processing.processors.base import Processor
+from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
+from data_designer.engine.registry.base import TaskRegistry
+
+
+class ProcessorRegistry(TaskRegistry[str, Processor, ConfigBase]): ...
+
+
+def create_default_processor_registry() -> ProcessorRegistry:
+    registry = ProcessorRegistry()
+    registry.register(ProcessorType.DROP_COLUMNS, DropColumnsProcessor, DropColumnsProcessorConfig, False)
+    return registry

data_designer/engine/processing/utils.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import logging
+from typing import Any, TypeVar, Union, overload
+
+import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+def concat_datasets(datasets: list[pd.DataFrame]) -> pd.DataFrame:
+    _verify_columns_are_unique(datasets)
+    _verify_dataset_lengths_are_equal(datasets)
+    emoji = " + ".join(["💾"] * len(datasets))
+    logger.info(f"({emoji}) Concatenating {len(datasets)} datasets")
+    return pd.concat([df for df in datasets], axis=1)
+
+
+# Overloads to help static type checker better understand
+# the input/output types of the deserialize_json_values function.
+@overload
+def deserialize_json_values(data: str) -> Union[dict[str, Any], list[Any], Any]: ...
+
+
+@overload
+def deserialize_json_values(data: list[T]) -> list[Any]: ...
+
+
+@overload
+def deserialize_json_values(data: dict[K, V]) -> dict[K, Any]: ...
+
+
+@overload
+def deserialize_json_values(data: T) -> T: ...
+
+
+def deserialize_json_values(data):
+    """De-serialize JSON strings in various input formats.
+
+    Args:
+        data: Input data in one of four formats:
+            - Single string (JSON string to deserialize)
+            - List of strings (list of JSON strings to deserialize)
+            - Dictionary (potentially with nested JSON strings to deserialize)
+            - Some other object that can't be deserialized.
+
+
+    Returns:
+        Deserialized data in the corresponding format:
+            - Dictionary (when input is a single string)
+            - List of dictionaries (when input is a list of strings)
+            - Dictionary (when input is a dictionary, with nested JSON strings deserialized)
+            - The original object (if there is no deserialization to perform)
+    """
+    # Case 1: Single string input
+    if isinstance(data, str):
+        try:
+            return json.loads(data)
+        except json.JSONDecodeError:
+            return data
+
+    # Case 2: List of strings input
+    elif isinstance(data, list):
+        result = []
+        for item in data:
+            if isinstance(item, str):
+                try:
+                    result.append(json.loads(item))
+                except json.JSONDecodeError:
+                    result.append(item)
+            else:
+                # If list contains non-string items, recursively process them
+                result.append(deserialize_json_values(item))
+        return result
+
+    # Case 3: Dictionary input with potential nested JSON strings
+    elif isinstance(data, dict):
+        result = {}
+        for key, value in data.items():
+            if isinstance(value, str):
+                try:
+                    result[key] = json.loads(value)
+                except json.JSONDecodeError:
+                    result[key] = value
+            elif isinstance(value, (dict, list)):
+                # Recursively process nested dictionaries and lists
+                result[key] = deserialize_json_values(value)
+            else:
+                result[key] = value
+        return result
+
+    # Fallback for other data types
+    else:
+        return data
+
+
+def _verify_columns_are_unique(datasets: list[pd.DataFrame]) -> None:
+    joined_columns = set()
+    for df in datasets:
+        columns = set(df.columns)
+        overlapping_columns = joined_columns & columns
+        if len(overlapping_columns) > 0:
+            raise ValueError(
+                f"🛑 Input datasets have overlapping columns: {overlapping_columns} "
+                "Please ensure that the column names are unique."
+            )
+        joined_columns.update(columns)
+
+
+def _verify_dataset_lengths_are_equal(datasets: list[pd.DataFrame]) -> None:
+    if len(set([len(df) for df in datasets])) > 1:
+        raise ValueError(
+            "🛑 Input datasets have different lengths. Please ensure that the datasets have the same number of rows."
+        )

data_designer/engine/registry/base.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import threading
+from typing import Any, Generic, Type, TypeVar
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.utils.type_helpers import StrEnum
+from data_designer.engine.configurable_task import ConfigurableTask
+from data_designer.engine.registry.errors import NotFoundInRegistryError, RegistryItemNotTypeError
+
+EnumNameT = TypeVar("EnumNameT", bound=StrEnum)
+TaskT = TypeVar("TaskT", bound=ConfigurableTask)
+TaskConfigT = TypeVar("TaskConfigT", bound=ConfigBase)
+
+
+class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
+    # registered type name -> type
+    _registry: dict[EnumNameT, Type[TaskT]] = {}
+    # type -> registered type name
+    _reverse_registry: dict[Type[TaskT], EnumNameT] = {}
+
+    # registered type name -> config type
+    _config_registry: dict[EnumNameT, Type[TaskConfigT]] = {}
+    # config type -> registered type name
+    _reverse_config_registry: dict[Type[TaskConfigT], EnumNameT] = {}
+
+    # all registries are singletons
+    _instance = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def register(
+        cls,
+        name: EnumNameT,
+        task: Type[TaskT],
+        config: Type[TaskConfigT],
+        raise_on_collision: bool = False,
+    ) -> None:
+        if cls._has_been_registered(name):
+            if not raise_on_collision:
+                return
+            raise ValueError(f"{name} has already been registered!")
+
+        cls._raise_if_not_type(task)
+        cls._raise_if_not_type(config)
+
+        with cls._lock:
+            cls._registry[name] = task
+            cls._reverse_registry[task] = name
+            cls._config_registry[name] = config
+            cls._reverse_config_registry[config] = name
+
+    @classmethod
+    def get_task_type(cls, name: EnumNameT) -> Type[TaskT]:
+        cls._raise_if_not_registered(name, cls._registry)
+        return cls._registry[name]
+
+    @classmethod
+    def get_config_type(cls, name: EnumNameT) -> Type[TaskConfigT]:
+        cls._raise_if_not_registered(name, cls._config_registry)
+        return cls._config_registry[name]
+
+    @classmethod
+    def get_registered_name(cls, task: Type[TaskT]) -> EnumNameT:
+        cls._raise_if_not_registered(task, cls._reverse_registry)
+        return cls._reverse_registry[task]
+
+    @classmethod
+    def get_for_config_type(cls, config: Type[TaskConfigT]) -> Type[TaskT]:
+        cls._raise_if_not_registered(config, cls._reverse_config_registry)
+        name = cls._reverse_config_registry[config]
+        return cls.get_task_type(name)
+
+    @classmethod
+    def _has_been_registered(cls, name: EnumNameT) -> bool:
+        return name in cls._registry
+
+    @classmethod
+    def _raise_if_not_registered(cls, key: EnumNameT | Type[TaskT] | Type[TaskConfigT], mapping: dict) -> None:
+        if not (isinstance(key, StrEnum) or isinstance(key, str)):
+            cls._raise_if_not_type(key)
+        if key not in mapping:
+            raise NotFoundInRegistryError(f"{key} not found in registry")
+
+    @classmethod
+    def _raise_if_not_type(cls, obj: Any) -> None:
+        if not isinstance(obj, type):
+            raise RegistryItemNotTypeError(f"{obj} is not a class!")
+
+    def __new__(cls, *args, **kwargs):
+        """Registry is a singleton."""
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super().__new__(cls)
+        return cls._instance

data_designer/engine/registry/data_designer_registry.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.engine.analysis.column_profilers.registry import (
+    ColumnProfilerRegistry,
+    create_default_column_profiler_registry,
+)
+from data_designer.engine.column_generators.registry import (
+    ColumnGeneratorRegistry,
+    create_default_column_generator_registry,
+)
+from data_designer.engine.processing.processors.registry import ProcessorRegistry, create_default_processor_registry
+
+
+class DataDesignerRegistry:
+    def __init__(
+        self,
+        *,
+        column_generator_registry: ColumnGeneratorRegistry | None = None,
+        column_profiler_registry: ColumnProfilerRegistry | None = None,
+        processor_registry: ProcessorRegistry | None = None,
+    ):
+        self._column_generator_registry = column_generator_registry or create_default_column_generator_registry()
+        self._column_profiler_registry = column_profiler_registry or create_default_column_profiler_registry()
+        self._processor_registry = processor_registry or create_default_processor_registry()
+
+    @property
+    def column_generators(self) -> ColumnGeneratorRegistry:
+        return self._column_generator_registry
+
+    @property
+    def column_profilers(self) -> ColumnProfilerRegistry:
+        return self._column_profiler_registry
+
+    @property
+    def processors(self) -> ProcessorRegistry:
+        return self._processor_registry

data_designer/engine/registry/errors.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.engine.errors import DataDesignerError
+
+
+class NotFoundInRegistryError(DataDesignerError): ...
+
+
+class RegistryItemNotTypeError(DataDesignerError): ...

data_designer/engine/resources/managed_dataset_generator.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+import pandas as pd
+
+from data_designer.engine.resources.managed_dataset_repository import ManagedDatasetRepository
+
+
+class ManagedDatasetGenerator:
+    def __init__(self, managed_datasets: ManagedDatasetRepository, dataset_name: str):
+        self.managed_datasets = managed_datasets
+        self.dataset_name = dataset_name
+
+    def generate_samples(
+        self,
+        size: int = 1,
+        evidence: dict[str, Any | list[Any]] = {},
+    ) -> pd.DataFrame:
+        parameters = []
+        query = f"select * from {self.dataset_name}"
+        if evidence:
+            where_conditions = []
+            for column, values in evidence.items():
+                if values:
+                    values = values if isinstance(values, list) else [values]
+                    formatted_values = ["?"] * len(values)
+                    condition = f"{column} IN ({', '.join(formatted_values)})"
+                    where_conditions.append(condition)
+                    parameters.extend(values)
+            if where_conditions:
+                query += " where " + " and ".join(where_conditions)
+        query += f" order by random() limit {size}"
+        return self.managed_datasets.query(query, parameters)