data-designer 0.1.0__py3-none-any.whl

data_designer/engine/column_generators/generators/validation.py

@@ -0,0 +1,147 @@
+# 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.column_configs import ValidationColumnConfig
+from data_designer.config.errors import InvalidConfigError
+from data_designer.config.utils.code_lang import SQL_DIALECTS, CodeLang
+from data_designer.config.validator_params import (
+    ValidatorParamsT,
+    ValidatorType,
+)
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    GeneratorMetadata,
+)
+from data_designer.engine.dataset_builders.utils.concurrency import ConcurrentThreadExecutor
+from data_designer.engine.errors import DataDesignerRuntimeError
+from data_designer.engine.validators import (
+    BaseValidator,
+    LocalCallableValidator,
+    PythonValidator,
+    RemoteValidator,
+    SQLValidator,
+    ValidationResult,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_validator_from_params(validator_type: ValidatorType, validator_params: ValidatorParamsT) -> BaseValidator:
+    if validator_type == ValidatorType.CODE:
+        if validator_params.code_lang == CodeLang.PYTHON:
+            return PythonValidator(validator_params)
+        elif validator_params.code_lang in SQL_DIALECTS:
+            return SQLValidator(validator_params)
+    elif validator_type == ValidatorType.REMOTE:
+        return RemoteValidator(validator_params)
+    else:
+        return LocalCallableValidator(validator_params)
+
+
+class ValidationColumnGenerator(ColumnGenerator[ValidationColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="validate",
+            description="Validate data.",
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            required_resources=None,
+        )
+
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame:
+        logger.info(f"🔍 Validating column {self.config.name!r} with {len(data)} records")
+        logger.info(f"  |-- target columns: {self.config.target_columns}")
+        logger.info(f"  |-- validator type: {self.config.validator_type}")
+        logger.info(f"  |-- validator params: {self.config.validator_params}")
+        logger.info(f"  |-- batch size: {self.config.batch_size}")
+
+        validator = get_validator_from_params(self.config.validator_type, self.config.validator_params)
+
+        # Check if the target columns are present in the dataset
+        missing_columns = set(self.config.target_columns) - set(data.columns)
+        if missing_columns:
+            raise InvalidConfigError(
+                f"Target columns {missing_columns} defined in validation column {self.config.name!r} are missing in dataset"
+            )
+
+        # Check whether to pass single columns or multiple columns to the validator
+        validate_columns_separately = False
+        if self.config.validator_type == ValidatorType.CODE and len(self.config.target_columns) > 1:
+            # Code validator expects single column input, so we validate each column separately
+            validate_columns_separately = True
+
+            columns_to_validate = [[col] for col in self.config.target_columns]
+        else:
+            columns_to_validate = [self.config.target_columns]
+
+        outputs_as_dicts = None
+        for cols in columns_to_validate:
+            # Filter the dataset to only include the target columns, and convert to a list of dictionaries
+            records = data[cols].to_dict(orient="records")
+
+            batched_records = [
+                records[batch_start : batch_start + self.config.batch_size]
+                for batch_start in range(0, len(records), self.config.batch_size)
+            ]
+
+            # Run validation in parallel or sequentially, depending on the validator type and parameters
+            if (
+                self.config.validator_type == ValidatorType.REMOTE
+                and self.config.validator_params.max_parallel_requests > 1
+            ):
+                concatenated_outputs = self._validate_in_parallel(validator, batched_records)
+            else:
+                concatenated_outputs = []
+                for batch in batched_records:
+                    concatenated_outputs.extend(self._validate_batch(validator, batch))
+
+            if validate_columns_separately:
+                if outputs_as_dicts is None:
+                    outputs_as_dicts = [{cols[0]: output.model_dump(mode="json")} for output in concatenated_outputs]
+                else:
+                    for dict_output in outputs_as_dicts:
+                        dict_output[cols[0]] = concatenated_outputs[0].model_dump(mode="json")
+            else:
+                outputs_as_dicts = [output.model_dump(mode="json") for output in concatenated_outputs]
+
+        validation_results = pd.DataFrame({self.config.name: outputs_as_dicts})
+        return pd.concat([data, validation_results], axis=1)
+
+    def _validate_in_parallel(self, validator: BaseValidator, batched_records: list[list[dict]]) -> pd.DataFrame:
+        """Run validation in parallel."""
+
+        outputs = [None] * len(batched_records)
+
+        def result_callback(result: ValidationResult, context: dict):
+            outputs[context["index"]] = result
+
+        def error_callback(error: Exception, context: dict):
+            outputs[context["index"]] = ValidationResult.empty(size=len(batched_records[context["index"]]))
+
+        with ConcurrentThreadExecutor(
+            max_workers=self.config.validator_params.max_parallel_requests,
+            column_name=self.config.name,
+            result_callback=result_callback,
+            error_callback=error_callback,
+        ) as executor:
+            for i, batch in enumerate(batched_records):
+                executor.submit(lambda batch: self._validate_batch(validator, batch), batch, context={"index": i})
+
+        if any(output is None for output in outputs):
+            raise DataDesignerRuntimeError("Validation task failed due to an unexpected error in parallel execution")
+
+        # Concatenate the outputs and convert to a DataFrame
+        return sum([output.data for output in outputs], [])
+
+    def _validate_batch(self, validator: BaseValidator, batch: list[dict]) -> ValidationResult:
+        try:
+            return validator.run_validation(batch)
+        except Exception as e:
+            error_to_display = str(e).replace("\n", "\n  ")  # add spaces to improve readability
+            logger.error(f"Batch could not be validated:\n  {error_to_display}")
+            raise e

data_designer/engine/column_generators/registry.py

@@ -0,0 +1,56 @@
+# 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.column_configs import (
+    ExpressionColumnConfig,
+    LLMCodeColumnConfig,
+    LLMJudgeColumnConfig,
+    LLMStructuredColumnConfig,
+    LLMTextColumnConfig,
+    ValidationColumnConfig,
+)
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
+from data_designer.engine.column_generators.generators.llm_generators import (
+    LLMCodeCellGenerator,
+    LLMJudgeCellGenerator,
+    LLMStructuredCellGenerator,
+    LLMTextCellGenerator,
+)
+from data_designer.engine.column_generators.generators.samplers import SamplerColumnGenerator
+from data_designer.engine.column_generators.generators.seed_dataset import SeedDatasetColumnGenerator
+from data_designer.engine.column_generators.generators.validation import ValidationColumnGenerator
+from data_designer.engine.dataset_builders.multi_column_configs import (
+    SamplerMultiColumnConfig,
+    SeedDatasetMultiColumnConfig,
+)
+from data_designer.engine.registry.base import TaskRegistry
+from data_designer.plugins.plugin import PluginType
+from data_designer.plugins.registry import PluginRegistry
+
+
+class ColumnGeneratorRegistry(TaskRegistry[DataDesignerColumnType, ColumnGenerator, ConfigBase]): ...
+
+
+def create_default_column_generator_registry(with_plugins: bool = True) -> ColumnGeneratorRegistry:
+    registry = ColumnGeneratorRegistry()
+    registry.register(DataDesignerColumnType.LLM_TEXT, LLMTextCellGenerator, LLMTextColumnConfig)
+    registry.register(DataDesignerColumnType.LLM_CODE, LLMCodeCellGenerator, LLMCodeColumnConfig)
+    registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig)
+    registry.register(DataDesignerColumnType.EXPRESSION, ExpressionColumnGenerator, ExpressionColumnConfig)
+    registry.register(DataDesignerColumnType.SAMPLER, SamplerColumnGenerator, SamplerMultiColumnConfig)
+    registry.register(DataDesignerColumnType.SEED_DATASET, SeedDatasetColumnGenerator, SeedDatasetMultiColumnConfig)
+    registry.register(DataDesignerColumnType.VALIDATION, ValidationColumnGenerator, ValidationColumnConfig)
+    registry.register(DataDesignerColumnType.LLM_STRUCTURED, LLMStructuredCellGenerator, LLMStructuredColumnConfig)
+
+    if with_plugins:
+        for plugin in PluginRegistry().get_plugins(PluginType.COLUMN_GENERATOR):
+            registry.register(
+                DataDesignerColumnType(plugin.name),
+                plugin.task_cls,
+                plugin.config_cls,
+            )
+
+    return registry

data_designer/engine/column_generators/utils/errors.py

@@ -0,0 +1,13 @@
+# 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 PromptTemplateRenderError(DataDesignerError): ...
+
+
+class ExpressionTemplateRenderError(DataDesignerError): ...
+
+
+class SeedDatasetError(DataDesignerError): ...

data_designer/engine/column_generators/utils/judge_score_factory.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from enum import Enum
+from typing import Type
+
+from pydantic import BaseModel, ConfigDict, Field, create_model
+
+from data_designer.config.column_configs import Score
+
+SCORING_FORMAT = "* {score}: {description}"
+SCORE_FIELD_DESCRIPTION_FORMAT = "Score Descriptions for {enum_name}:\n{scoring}"
+
+
+class BaseJudgeResponse(BaseModel):
+    """Base model for all rubrics."""
+
+    model_config = ConfigDict(use_enum_values=True)
+    reasoning: str = Field(..., description="Reasoning for the assigned score.")
+
+
+def _stringify_scoring(options: dict, enum_type: Type[Enum]) -> str:
+    """Convert score descriptions into a single text block."""
+    list_block = "\n".join(
+        [SCORING_FORMAT.format(score=score, description=description) for score, description in options.items()]
+    )
+    return SCORE_FIELD_DESCRIPTION_FORMAT.format(enum_name=enum_type.__name__, scoring=list_block)
+
+
+def create_judge_response_model(score: Score) -> Type[BaseJudgeResponse]:
+    """Create a JudgeResponse data type."""
+    enum_members = {}
+    for option in score.options.keys():
+        member_name = f"VALUE_{option}"
+        enum_members[member_name] = option
+
+    DynamicScaleEnum = Enum(f"{score.name}Enum", enum_members)
+    options = _stringify_scoring(score.options, enum_type=DynamicScaleEnum)
+
+    return create_model(
+        score.name,
+        __doc__=score.description if score.description else None,
+        __base__=BaseJudgeResponse,
+        score=(DynamicScaleEnum, Field(..., description=options)),
+    )
+
+
+def create_judge_structured_output_model(
+    judge_responses: list[Type[BaseJudgeResponse]],
+) -> Type[BaseModel]:
+    """Create a JudgeStructuredOutput class dynamically."""
+    return create_model(
+        "JudgeStructuredOutput",
+        __doc__=f"Response schema for scores with the following names: {[response.__name__ for response in judge_responses]}.",
+        __base__=BaseModel,
+        **{response.__name__.lower(): (response, ...) for response in judge_responses},
+    )

data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import logging
+
+from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.models import ModelConfig
+from data_designer.config.utils.code_lang import CodeLang
+from data_designer.config.utils.misc import get_prompt_template_keywords
+from data_designer.config.utils.type_helpers import StrEnum
+from data_designer.engine.column_generators.utils.errors import PromptTemplateRenderError
+from data_designer.engine.column_generators.utils.judge_score_factory import (
+    create_judge_response_model,
+    create_judge_structured_output_model,
+)
+from data_designer.engine.models.recipes.base import ResponseRecipe
+from data_designer.engine.models.recipes.response_recipes import (
+    CodeResponseRecipe,
+    PydanticResponseRecipe,
+    StructuredResponseRecipe,
+    TextResponseRecipe,
+)
+from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
+from data_designer.engine.processing.ginja.exceptions import UserTemplateError, UserTemplateUnsupportedFiltersError
+
+logger = logging.getLogger(__name__)
+
+
+class PromptType(StrEnum):
+    SYSTEM_PROMPT = "system_prompt"
+    USER_PROMPT = "user_prompt"
+
+
+class RecordBasedPromptRenderer(WithJinja2UserTemplateRendering):
+    def __init__(self, response_recipe: ResponseRecipe, *, error_message_context: dict[str, str] | None = None):
+        self.response_recipe = response_recipe
+        self._error_message_context = error_message_context
+
+    def render(self, *, prompt_template: str | None, record: dict, prompt_type: PromptType) -> str | None:
+        self._prepare_environment(prompt_template=prompt_template, record=record, prompt_type=prompt_type)
+        rendered_prompt = self.render_multi_template(prompt_type, record) if prompt_template else ""
+        recipe_applicator = (
+            self.response_recipe.apply_recipe_to_user_prompt
+            if prompt_type == PromptType.USER_PROMPT
+            else self.response_recipe.apply_recipe_to_system_prompt
+        )
+        return recipe_applicator(rendered_prompt)
+
+    def _prepare_environment(self, *, prompt_template: str | None, record: dict, prompt_type: PromptType) -> None:
+        try:
+            self.prepare_jinja2_multi_template_renderer(
+                template_name=prompt_type.value,
+                prompt_template=prompt_template,
+                dataset_variables=list(record.keys()),
+            )
+        except (UserTemplateUnsupportedFiltersError, UserTemplateError) as exc:
+            template_variables = get_prompt_template_keywords(prompt_template)
+            missing_columns = list(set(template_variables) - set(record.keys()))
+
+            error_msg = (
+                f"There was an error preparing the {prompt_type.value.replace('_', ' ')} "
+                "template. Please double check that the template is valid Jinja2 syntax, that all "
+                "referenced variables are defined, and that any filters you are using are supported."
+            )
+            if len(missing_columns) > 0:
+                error_msg += f"\nThe following {missing_columns} columns are missing!"
+            if self._error_message_context is not None:
+                error_msg += f"\n{json.dumps(self._error_message_context, indent=2)}"
+            logger.error(f"🛑 {error_msg}")
+            raise PromptTemplateRenderError(f"{exc!s} {error_msg}")
+
+
+def create_response_recipe(
+    column_config: SingleColumnConfig, model_config: ModelConfig | None = None
+) -> ResponseRecipe:
+    if model_config and column_config.model_alias != model_config.alias:
+        raise ValueError(
+            f"Column config model alias {column_config.model_alias} does not match model config alias {model_config.alias}"
+        )
+    if column_config.column_type == DataDesignerColumnType.LLM_TEXT:
+        return TextResponseRecipe()
+    if column_config.column_type == DataDesignerColumnType.LLM_CODE:
+        return CodeResponseRecipe(
+            syntax=CodeLang.parse_lang(column_config.code_lang),
+        )
+    if column_config.column_type == DataDesignerColumnType.LLM_STRUCTURED:
+        return StructuredResponseRecipe(
+            json_schema=column_config.output_format,
+        )
+    if column_config.column_type == DataDesignerColumnType.LLM_JUDGE:
+        return PydanticResponseRecipe(
+            data_type=create_judge_structured_output_model(
+                [create_judge_response_model(s) for s in column_config.scores]
+            ),
+        )
+    raise ValueError(f"No response recipe found for column type: {column_config.column_type}")

data_designer/engine/configurable_task.py

@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Generic, Type, TypeVar, get_origin
+
+import pandas as pd
+
+from data_designer.config.base import ConfigBase
+from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+from data_designer.engine.resources.resource_provider import ResourceProvider, ResourceType
+
+DataT = TypeVar("DataT", dict, pd.DataFrame)
+TaskConfigT = TypeVar("ConfigT", bound=ConfigBase)
+
+
+class ConfigurableTaskMetadata(ConfigBase):
+    name: str
+    description: str
+    required_resources: list[ResourceType] | None
+
+
+class ConfigurableTask(ABC, Generic[TaskConfigT]):
+    def __init__(self, config: TaskConfigT, *, resource_provider: ResourceProvider | None):
+        self._config = self.get_config_type().model_validate(config)
+        self._resource_provider = resource_provider
+        self._validate_resources()
+        self._validate()
+        self._initialize()
+
+    @classmethod
+    def get_config_type(cls) -> Type[TaskConfigT]:
+        for base in cls.__orig_bases__:
+            if hasattr(base, "__args__") and len(base.__args__) == 1:
+                arg = base.__args__[0]
+                origin = get_origin(arg) or arg
+                if isinstance(origin, type) and issubclass(origin, ConfigBase):
+                    return base.__args__[0]
+        raise TypeError(
+            f"Could not determine config type for `{cls.__name__}`. Please ensure that the "
+            "`ConfigurableTask` is defined with a generic type argument, where the type argument "
+            "is a subclass of `ConfigBase`."
+        )
+
+    @property
+    def artifact_path(self) -> Path:
+        return self.artifact_storage.artifact_path
+
+    @property
+    def artifact_storage(self) -> ArtifactStorage:
+        return self.resource_provider.artifact_storage
+
+    @property
+    def base_dataset_path(self) -> Path:
+        return self.artifact_storage.base_dataset_path
+
+    @property
+    def config(self) -> TaskConfigT:
+        return self._config
+
+    @property
+    def resource_provider(self) -> ResourceProvider:
+        if self._resource_provider is None:
+            raise ValueError(f"No resource provider provided for the `{self.metadata().name}` task.")
+        return self._resource_provider
+
+    @staticmethod
+    @abstractmethod
+    def metadata() -> ConfigurableTaskMetadata: ...
+
+    def _initialize(self) -> None:
+        """An internal method for custom initialization logic, which will be called in the constructor."""
+
+    def _validate(self) -> None:
+        """An internal method for custom validation logic, which will be called in the constructor."""
+
+    def _validate_resources(self) -> None:
+        for resource in self.metadata().required_resources or []:
+            if resource is not None:
+                if getattr(self.resource_provider, ResourceType(resource).value) is None:
+                    raise ValueError(f"Resource {resource} is required for the `{self.metadata().name}`")

data_designer/engine/dataset_builders/artifact_storage.py

@@ -0,0 +1,181 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import logging
+from pathlib import Path
+import shutil
+from typing import Union
+
+import pandas as pd
+from pydantic import BaseModel, field_validator, model_validator
+
+from data_designer.config.utils.io_helpers import read_parquet_dataset
+from data_designer.config.utils.type_helpers import StrEnum, resolve_string_enum
+from data_designer.engine.dataset_builders.errors import ArtifactStorageError
+
+logger = logging.getLogger(__name__)
+
+BATCH_FILE_NAME_FORMAT = "batch_{batch_number:05d}.parquet"
+
+
+class BatchStage(StrEnum):
+    PARTIAL_RESULT = "partial_results_path"
+    FINAL_RESULT = "final_dataset_path"
+    DROPPED_COLUMNS = "dropped_columns_dataset_path"
+
+
+class ArtifactStorage(BaseModel):
+    artifact_path: Path | str
+    dataset_name: str = "dataset"
+    final_dataset_folder_name: str = "parquet-files"
+    partial_results_folder_name: str = "tmp-partial-parquet-files"
+    dropped_columns_folder_name: str = "dropped-columns-parquet-files"
+
+    @property
+    def artifact_path_exists(self) -> bool:
+        return self.artifact_path.exists()
+
+    @property
+    def base_dataset_path(self) -> Path:
+        return self.artifact_path / self.dataset_name
+
+    @property
+    def dropped_columns_dataset_path(self) -> Path:
+        return self.base_dataset_path / self.dropped_columns_folder_name
+
+    @property
+    def final_dataset_path(self) -> Path:
+        return self.base_dataset_path / self.final_dataset_folder_name
+
+    @property
+    def metadata_file_path(self) -> Path:
+        return self.base_dataset_path / "metadata.json"
+
+    @property
+    def partial_results_path(self) -> Path:
+        return self.base_dataset_path / self.partial_results_folder_name
+
+    @field_validator("artifact_path")
+    def validate_artifact_path(cls, v: Union[Path, str]) -> Path:
+        v = Path(v)
+        if not v.is_dir():
+            raise ArtifactStorageError("Artifact path must exist and be a directory")
+        return v
+
+    @model_validator(mode="after")
+    def validate_folder_names(self):
+        folder_names = [
+            self.dataset_name,
+            self.final_dataset_folder_name,
+            self.partial_results_folder_name,
+            self.dropped_columns_folder_name,
+        ]
+
+        for name in folder_names:
+            if len(name) == 0:
+                raise ArtifactStorageError("🛑 Directory names must be non-empty strings.")
+
+        if len(set(folder_names)) != len(folder_names):
+            raise ArtifactStorageError("🛑 Folder names must be unique (no collisions allowed).")
+
+        invalid_chars = {"<", ">", ":", '"', "/", "\\", "|", "?", "*"}
+        for name in folder_names:
+            if any(char in invalid_chars for char in name):
+                raise ArtifactStorageError(f"🛑 Directory name '{name}' contains invalid characters.")
+
+        return self
+
+    @staticmethod
+    def mkdir_if_needed(path: Path | str) -> Path:
+        """Create the directory if it does not exist."""
+        path = Path(path)
+        if not path.exists():
+            logger.debug(f"📁 Creating directory: {path}")
+            path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    @staticmethod
+    def read_parquet_files(path: Path) -> pd.DataFrame:
+        return read_parquet_dataset(path)
+
+    def create_batch_file_path(
+        self,
+        batch_number: int,
+        batch_stage: BatchStage,
+    ) -> Path:
+        if batch_number < 0:
+            raise ArtifactStorageError("🛑 Batch number must be non-negative.")
+        return self._get_stage_path(batch_stage) / BATCH_FILE_NAME_FORMAT.format(batch_number=batch_number)
+
+    def load_dataset(self, batch_stage: BatchStage = BatchStage.FINAL_RESULT) -> pd.DataFrame:
+        return read_parquet_dataset(self._get_stage_path(batch_stage))
+
+    def load_dataset_with_dropped_columns(self) -> pd.DataFrame:
+        # The pyarrow backend has better support for nested data types.
+        df = self.load_dataset()
+        if (
+            self.dropped_columns_dataset_path.exists()
+            and self.create_batch_file_path(0, BatchStage.DROPPED_COLUMNS).is_file()
+        ):
+            logger.debug("Concatenating dropped columns to the final dataset.")
+            df_dropped = self.load_dataset(batch_stage=BatchStage.DROPPED_COLUMNS)
+            if len(df_dropped) != len(df):
+                raise ArtifactStorageError(
+                    "🛑 The dropped-columns dataset has a different number of rows than the main dataset. "
+                    "Something unexpected must have happened to the dataset builder's artifacts."
+                )
+            # To ensure indexes are aligned and avoid silent misalignment (which would introduce NaNs),
+            # check that the indexes are identical before concatenation.
+            if not df.index.equals(df_dropped.index):
+                raise ArtifactStorageError(
+                    "🛑 The indexes of the main and dropped columns DataFrames are not aligned. "
+                    "Something unexpected must have happened to the dataset builder's artifacts."
+                )
+            df = pd.concat([df, df_dropped], axis=1)
+        return df
+
+    def move_partial_result_to_final_file_path(self, batch_number: int) -> Path:
+        partial_result_path = self.create_batch_file_path(batch_number, batch_stage=BatchStage.PARTIAL_RESULT)
+        if not partial_result_path.exists():
+            raise ArtifactStorageError("🛑 Partial result file not found.")
+        self.mkdir_if_needed(self._get_stage_path(BatchStage.FINAL_RESULT))
+        final_file_path = self.create_batch_file_path(batch_number, batch_stage=BatchStage.FINAL_RESULT)
+        shutil.move(partial_result_path, final_file_path)
+        return final_file_path
+
+    def write_configs(self, json_file_name: str, configs: list[dict]) -> Path:
+        self.mkdir_if_needed(self.base_dataset_path)
+        with open(self.base_dataset_path / json_file_name, "w") as file:
+            json.dump([c.model_dump(mode="json") for c in configs], file, indent=4)
+        return self.base_dataset_path / json_file_name
+
+    def write_batch_to_parquet_file(
+        self,
+        batch_number: int,
+        dataframe: pd.DataFrame,
+        batch_stage: BatchStage,
+    ) -> Path:
+        file_path = self.create_batch_file_path(batch_number, batch_stage=batch_stage)
+        self.write_parquet_file(file_path.name, dataframe, batch_stage)
+        return file_path
+
+    def write_parquet_file(
+        self,
+        parquet_file_name: str,
+        dataframe: pd.DataFrame,
+        batch_stage: BatchStage,
+    ) -> Path:
+        self.mkdir_if_needed(self._get_stage_path(batch_stage))
+        file_path = self._get_stage_path(batch_stage) / parquet_file_name
+        dataframe.to_parquet(file_path, index=False)
+        return file_path
+
+    def write_metadata(self, metadata: dict) -> Path:
+        self.mkdir_if_needed(self.base_dataset_path)
+        with open(self.metadata_file_path, "w") as file:
+            json.dump(metadata, file)
+        return self.metadata_file_path
+
+    def _get_stage_path(self, stage: BatchStage) -> Path:
+        return getattr(self, resolve_string_enum(stage, BatchStage).value)