data-designer 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import json
 import logging
 
@@ -8,7 +10,7 @@ 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.misc import extract_keywords_from_jinja2_template
 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 (
@@ -56,7 +58,7 @@ class RecordBasedPromptRenderer(WithJinja2UserTemplateRendering):
                 dataset_variables=list(record.keys()),
             )
         except (UserTemplateUnsupportedFiltersError, UserTemplateError) as exc:
-            template_variables = get_prompt_template_keywords(prompt_template)
+            template_variables = extract_keywords_from_jinja2_template(prompt_template)
             missing_columns = list(set(template_variables) - set(record.keys()))
 
             error_msg = (

data_designer/engine/compiler.py

@@ -1,10 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 
 from data_designer.config.column_configs import SeedDatasetColumnConfig
-from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.errors import InvalidConfigError
 from data_designer.engine.resources.resource_provider import ResourceProvider
@@ -14,13 +15,9 @@ from data_designer.engine.validation import ViolationLevel, rich_print_violation
 logger = logging.getLogger(__name__)
 
 
-def compile_data_designer_config(
-    config_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider
-) -> DataDesignerConfig:
-    config = config_builder.build()
+def compile_data_designer_config(config: DataDesignerConfig, resource_provider: ResourceProvider) -> DataDesignerConfig:
     _resolve_and_add_seed_columns(config, resource_provider.seed_reader)
     _validate(config)
-
     return config
 
 

data_designer/engine/configurable_task.py

@@ -1,25 +1,24 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 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, TypeVar, get_origin
+from __future__ import annotations
 
-import pandas as pd
+from abc import ABC
+from pathlib import Path
+from typing import TYPE_CHECKING, Generic, TypeVar, get_origin
 
 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
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 DataT = TypeVar("DataT", dict, pd.DataFrame)
 TaskConfigT = TypeVar("ConfigT", bound=ConfigBase)
 
 
-class ConfigurableTaskMetadata(ConfigBase):
-    name: str
-    description: str
-
-
 class ConfigurableTask(ABC, Generic[TaskConfigT]):
     def __init__(self, config: TaskConfigT, resource_provider: ResourceProvider):
         self._config = self.get_config_type().model_validate(config)
@@ -57,14 +56,14 @@ class ConfigurableTask(ABC, Generic[TaskConfigT]):
     def config(self) -> TaskConfigT:
         return self._config
 
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__
+
     @property
     def resource_provider(self) -> ResourceProvider:
         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."""
 

data_designer/engine/dataset_builders/artifact_storage.py

@@ -1,23 +1,30 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import json
 import logging
 import shutil
 from datetime import datetime
 from functools import cached_property
 from pathlib import Path
+from typing import TYPE_CHECKING
 
-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
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
 BATCH_FILE_NAME_FORMAT = "batch_{batch_number:05d}.parquet"
+SDG_CONFIG_FILENAME = "sdg.json"
 
 
 class BatchStage(StrEnum):
@@ -164,12 +171,6 @@ class ArtifactStorage(BaseModel):
         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,
@@ -194,11 +195,89 @@ class ArtifactStorage(BaseModel):
         dataframe.to_parquet(file_path, index=False)
         return file_path
 
+    def get_parquet_file_paths(self) -> list[str]:
+        """Get list of parquet file paths relative to base_dataset_path.
+
+        Returns:
+            List of relative paths to parquet files in the final dataset folder.
+        """
+        return [str(f.relative_to(self.base_dataset_path)) for f in sorted(self.final_dataset_path.glob("*.parquet"))]
+
+    def get_processor_file_paths(self) -> dict[str, list[str]]:
+        """Get processor output files organized by processor name.
+
+        Returns:
+            Dictionary mapping processor names to lists of relative file paths.
+        """
+        processor_files: dict[str, list[str]] = {}
+        if self.processors_outputs_path.exists():
+            for processor_dir in sorted(self.processors_outputs_path.iterdir()):
+                if processor_dir.is_dir():
+                    processor_name = processor_dir.name
+                    processor_files[processor_name] = [
+                        str(f.relative_to(self.base_dataset_path))
+                        for f in sorted(processor_dir.rglob("*"))
+                        if f.is_file()
+                    ]
+        return processor_files
+
+    def get_file_paths(self) -> dict[str, list[str] | dict[str, list[str]]]:
+        """Get all file paths organized by type.
+
+        Returns:
+            Dictionary with 'parquet-files' and 'processor-files' keys.
+        """
+        file_paths = {
+            "parquet-files": self.get_parquet_file_paths(),
+        }
+        processor_file_paths = self.get_processor_file_paths()
+        if processor_file_paths:
+            file_paths["processor-files"] = processor_file_paths
+
+        return file_paths
+
+    def read_metadata(self) -> dict:
+        """Read metadata from the metadata.json file.
+
+        Returns:
+            Dictionary containing the metadata.
+
+        Raises:
+            FileNotFoundError: If metadata file doesn't exist.
+        """
+        with open(self.metadata_file_path, "r") as file:
+            return json.load(file)
+
     def write_metadata(self, metadata: dict) -> Path:
+        """Write metadata to the metadata.json file.
+
+        Args:
+            metadata: Dictionary containing metadata to write.
+
+        Returns:
+            Path to the written metadata file.
+        """
         self.mkdir_if_needed(self.base_dataset_path)
         with open(self.metadata_file_path, "w") as file:
-            json.dump(metadata, file)
+            json.dump(metadata, file, indent=4, sort_keys=True)
         return self.metadata_file_path
 
+    def update_metadata(self, updates: dict) -> Path:
+        """Update existing metadata with new fields.
+
+        Args:
+            updates: Dictionary of fields to add/update in metadata.
+
+        Returns:
+            Path to the updated metadata file.
+        """
+        try:
+            existing_metadata = self.read_metadata()
+        except FileNotFoundError:
+            existing_metadata = {}
+
+        existing_metadata.update(updates)
+        return self.write_metadata(existing_metadata)
+
     def _get_stage_path(self, stage: BatchStage) -> Path:
         return getattr(self, resolve_string_enum(stage, BatchStage).value)

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -12,9 +12,9 @@ import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
 
-import pandas as pd
-
 from data_designer.config.column_types import ColumnConfigT
+from data_designer.config.config_builder import BuilderConfig
+from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
@@ -27,40 +27,38 @@ from data_designer.engine.column_generators.generators.base import (
     GenerationStrategy,
 )
 from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
-from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+from data_designer.engine.compiler import compile_data_designer_config
+from data_designer.engine.dataset_builders.artifact_storage import SDG_CONFIG_FILENAME, ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
-from data_designer.engine.dataset_builders.multi_column_configs import (
-    DatasetBuilderColumnConfigT,
-    MultiColumnConfig,
-)
+from data_designer.engine.dataset_builders.multi_column_configs import MultiColumnConfig
 from data_designer.engine.dataset_builders.utils.concurrency import (
     MAX_CONCURRENCY_PER_NON_LLM_GENERATOR,
     ConcurrentThreadExecutor,
 )
-from data_designer.engine.dataset_builders.utils.dataset_batch_manager import (
-    DatasetBatchManager,
-)
+from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
+from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
 from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
 from data_designer.engine.resources.resource_provider import ResourceProvider
+from data_designer.lazy_heavy_imports import pd
 
 if TYPE_CHECKING:
+    import pandas as pd
+
     from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModelRegistry
     from data_designer.engine.models.usage import ModelUsageStats
 
 logger = logging.getLogger(__name__)
 
-
 _CLIENT_VERSION: str = importlib.metadata.version("data_designer")
 
 
 class ColumnWiseDatasetBuilder:
     def __init__(
         self,
-        column_configs: list[DatasetBuilderColumnConfigT],
-        processor_configs: list[ProcessorConfig],
+        data_designer_config: DataDesignerConfig,
         resource_provider: ResourceProvider,
         registry: DataDesignerRegistry | None = None,
     ):
@@ -68,8 +66,12 @@ class ColumnWiseDatasetBuilder:
         self._resource_provider = resource_provider
         self._records_to_drop: set[int] = set()
         self._registry = registry or DataDesignerRegistry()
-        self._column_configs = column_configs
-        self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(processor_configs)
+
+        self._data_designer_config = compile_data_designer_config(data_designer_config, resource_provider)
+        self._column_configs = compile_dataset_builder_column_configs(self._data_designer_config)
+        self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(
+            self._data_designer_config.processors or []
+        )
         self._validate_column_configs()
 
     @property
@@ -94,16 +96,15 @@ class ColumnWiseDatasetBuilder:
         self,
         *,
         num_records: int,
-        buffer_size: int,
         on_batch_complete: Callable[[Path], None] | None = None,
     ) -> Path:
-        self._write_configs()
         self._run_model_health_check_if_needed()
-
+        self._write_builder_config()
         generators = self._initialize_generators()
         start_time = time.perf_counter()
         group_id = uuid.uuid4().hex
 
+        buffer_size = self._resource_provider.run_config.buffer_size
         self.batch_manager.start(num_records=num_records, buffer_size=buffer_size)
         for batch_idx in range(self.batch_manager.num_batches):
             logger.info(f"⏳ Processing batch {batch_idx + 1} of {self.batch_manager.num_batches}")
@@ -157,6 +158,12 @@ class ColumnWiseDatasetBuilder:
             for config in self._column_configs
         ]
 
+    def _write_builder_config(self) -> None:
+        self.artifact_storage.mkdir_if_needed(self.artifact_storage.base_dataset_path)
+        BuilderConfig(data_designer=self._data_designer_config).to_json(
+            self.artifact_storage.base_dataset_path / SDG_CONFIG_FILENAME
+        )
+
     def _run_batch(
         self, generators: list[ColumnGenerator], *, batch_mode: str, save_partial_results: bool = True, group_id: str
     ) -> None:
@@ -164,15 +171,16 @@ class ColumnWiseDatasetBuilder:
         for generator in generators:
             generator.log_pre_generation()
             try:
+                generation_strategy = generator.get_generation_strategy()
                 if generator.can_generate_from_scratch and self.batch_manager.buffer_is_empty:
                     self._run_from_scratch_column_generator(generator)
-                elif generator.generation_strategy == GenerationStrategy.CELL_BY_CELL:
+                elif generation_strategy == GenerationStrategy.CELL_BY_CELL:
                     self._run_cell_by_cell_generator(generator)
-                elif generator.generation_strategy == GenerationStrategy.FULL_COLUMN:
+                elif generation_strategy == GenerationStrategy.FULL_COLUMN:
                     self._run_full_column_generator(generator)
                 else:
-                    logger.error(f"❌ Unknown generation strategy: {generator.generation_strategy}")
-                    raise DatasetGenerationError(f"🛑 Unknown generation strategy: {generator.generation_strategy}")
+                    logger.error(f"❌ Unknown generation strategy: {generation_strategy}")
+                    raise DatasetGenerationError(f"🛑 Unknown generation strategy: {generation_strategy}")
                 if save_partial_results:
                     self.batch_manager.write()
             except Exception as e:
@@ -210,9 +218,9 @@ class ColumnWiseDatasetBuilder:
             )
 
     def _fan_out_with_threads(self, generator: ColumnGeneratorWithModelRegistry, max_workers: int) -> None:
-        if generator.generation_strategy != GenerationStrategy.CELL_BY_CELL:
+        if generator.get_generation_strategy() != GenerationStrategy.CELL_BY_CELL:
             raise DatasetGenerationError(
-                f"Generator {generator.metadata().name} is not a {GenerationStrategy.CELL_BY_CELL} "
+                f"Generator {generator.name} is not a {GenerationStrategy.CELL_BY_CELL} "
                 "generator so concurrency through threads is not supported."
             )
 
@@ -228,6 +236,7 @@ class ColumnWiseDatasetBuilder:
             error_callback=self._worker_error_callback,
             shutdown_error_rate=settings.shutdown_error_rate,
             shutdown_error_window=settings.shutdown_error_window,
+            disable_early_shutdown=settings.disable_early_shutdown,
         ) as executor:
             for i, record in self.batch_manager.iter_current_batch():
                 executor.submit(lambda record: generator.generate(record), record, context={"index": i})
@@ -291,7 +300,7 @@ class ColumnWiseDatasetBuilder:
                 dataframe = processor.process(dataframe, current_batch_number=current_batch_number)
             except Exception as e:
                 raise DatasetProcessingError(
-                    f"🛑 Failed to process dataset with processor {processor.metadata().name} in stage {stage}: {e}"
+                    f"🛑 Failed to process dataset with processor {processor.name} in stage {stage}: {e}"
                 ) from e
         return dataframe
 
@@ -306,16 +315,6 @@ class ColumnWiseDatasetBuilder:
     def _worker_result_callback(self, result: dict, *, context: dict | None = None) -> None:
         self.batch_manager.update_record(context["index"], result)
 
-    def _write_configs(self) -> None:
-        self.artifact_storage.write_configs(
-            json_file_name="column_configs.json",
-            configs=self._column_configs,
-        )
-        self.artifact_storage.write_configs(
-            json_file_name="model_configs.json",
-            configs=self._resource_provider.model_registry.model_configs.values(),
-        )
-
     def _emit_batch_inference_events(
         self, batch_mode: str, usage_deltas: dict[str, ModelUsageStats], group_id: str
     ) -> None:

data_designer/engine/dataset_builders/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.engine.errors import DataDesignerError
 
 

data_designer/engine/dataset_builders/multi_column_configs.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from abc import ABC
 from typing import TypeAlias
 

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -96,6 +96,7 @@ class ConcurrentThreadExecutor:
         error_callback: ErrorCallbackWithContext | None = None,
         shutdown_error_rate: float = 0.50,
         shutdown_error_window: int = 10,
+        disable_early_shutdown: bool = False,
     ):
         self._executor = None
         self._column_name = column_name
@@ -106,6 +107,7 @@ class ConcurrentThreadExecutor:
         self._error_callback = error_callback
         self._shutdown_error_rate = shutdown_error_rate
         self._shutdown_window_size = shutdown_error_window
+        self._disable_early_shutdown = disable_early_shutdown
         self._results = ExecutorResults(failure_threshold=shutdown_error_rate)
 
     @property
@@ -139,7 +141,7 @@ class ConcurrentThreadExecutor:
 
     def __exit__(self, exc_type, exc_value, traceback):
         self._shutdown_executor()
-        if self._results.early_shutdown is True:
+        if not self._disable_early_shutdown and self._results.early_shutdown is True:
             self._raise_task_error()
 
     def _shutdown_executor(self) -> None:
@@ -160,7 +162,7 @@ class ConcurrentThreadExecutor:
         if self._executor is None:
             raise RuntimeError("Executor is not initialized, this class should be used as a context manager.")
 
-        if self._results.early_shutdown:
+        if not self._disable_early_shutdown and self._results.early_shutdown:
             self._shutdown_executor()
             self._raise_task_error()
 
@@ -176,7 +178,9 @@ class ConcurrentThreadExecutor:
                 with self._lock:
                     self._results.completed_count += 1
                     self._results.error_trap.handle_error(err)
-                    if self._results.is_error_rate_exceeded(self._shutdown_window_size):
+                    if not self._disable_early_shutdown and self._results.is_error_rate_exceeded(
+                        self._shutdown_window_size
+                    ):
                         # Signal to shutdown early on the next submission (if received).
                         # We cannot trigger shutdown from within this thread as it can
                         # cause a deadlock.
@@ -196,7 +200,12 @@ class ConcurrentThreadExecutor:
             # We'll re-raise a custom error that can be handled at the call-site and the summary
             # can also be inspected.
             self._semaphore.release()
-            if not isinstance(err, RuntimeError) and "after shutdown" not in str(err):
+            is_shutdown_error = isinstance(err, RuntimeError) and (
+                "after shutdown" in str(err) or "Pool shutdown" in str(err)
+            )
+            if not is_shutdown_error:
+                raise err
+            if self._disable_early_shutdown:
                 raise err
             self._raise_task_error()
 

data_designer/engine/dataset_builders/utils/config_compiler.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.processors import ProcessorConfig

data_designer/engine/dataset_builders/utils/dag.py

@@ -1,13 +1,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import logging
+from __future__ import annotations
 
-import networkx as nx
+import logging
+from typing import TYPE_CHECKING
 
 from data_designer.config.column_types import ColumnConfigT
 from data_designer.engine.column_generators.utils.generator_classification import column_type_used_in_execution_dag
 from data_designer.engine.dataset_builders.utils.errors import DAGCircularDependencyError
+from data_designer.lazy_heavy_imports import nx
+
+if TYPE_CHECKING:
+    import networkx as nx
 
 logger = logging.getLogger(__name__)
 

data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

@@ -1,16 +1,20 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 import shutil
 from pathlib import Path
-from typing import Callable, Container, Iterator
-
-import pandas as pd
-import pyarrow.parquet as pq
+from typing import TYPE_CHECKING, Callable, Container, Iterator
 
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage, BatchStage
 from data_designer.engine.dataset_builders.utils.errors import DatasetBatchManagementError
+from data_designer.lazy_heavy_imports import pd, pq
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow.parquet as pq
 
 logger = logging.getLogger(__name__)
 
@@ -69,7 +73,7 @@ class DatasetBatchManager:
     def drop_records(self, index: Container[int]) -> None:
         self._buffer = [record for i, record in enumerate(self._buffer) if i not in index]
 
-    def finish_batch(self, on_complete: Callable[[Path], None] | None = None) -> Path:
+    def finish_batch(self, on_complete: Callable[[Path], None] | None = None) -> Path | None:
         """Finish the batch by moving the results from the partial results path to the final parquet folder.
 
         Returns:
@@ -78,29 +82,35 @@ class DatasetBatchManager:
         if self._current_batch_number >= self.num_batches:
             raise DatasetBatchManagementError("🛑 All batches have been processed.")
 
-        if not self.write():
-            raise DatasetBatchManagementError("🛑 Batch finished without any results to write.")
-
-        final_file_path = self.artifact_storage.move_partial_result_to_final_file_path(self._current_batch_number)
-
-        self.artifact_storage.write_metadata(
-            {
-                "target_num_records": sum(self.num_records_list),
-                "total_num_batches": self.num_batches,
-                "buffer_size": self._buffer_size,
-                "schema": {field.name: str(field.type) for field in pq.read_schema(final_file_path)},
-                "file_paths": [str(f) for f in sorted(self.artifact_storage.final_dataset_path.glob("*.parquet"))],
-                "num_records": self.num_records_list[: self._current_batch_number + 1],
-                "num_completed_batches": self._current_batch_number + 1,
-                "dataset_name": self.artifact_storage.dataset_name,
-            }
-        )
+        if self.write() is not None:
+            final_file_path = self.artifact_storage.move_partial_result_to_final_file_path(self._current_batch_number)
+
+            self.artifact_storage.write_metadata(
+                {
+                    "target_num_records": sum(self.num_records_list),
+                    "total_num_batches": self.num_batches,
+                    "buffer_size": self._buffer_size,
+                    "schema": {field.name: str(field.type) for field in pq.read_schema(final_file_path)},
+                    "file_paths": self.artifact_storage.get_file_paths(),
+                    "num_completed_batches": self._current_batch_number + 1,
+                    "dataset_name": self.artifact_storage.dataset_name,
+                }
+            )
+
+            if on_complete:
+                on_complete(final_file_path)
+        else:
+            final_file_path = None
+
+            logger.warning(
+                f"⚠️ Batch {self._current_batch_number + 1} finished without any results to write. "
+                "A partial dataset containing the currently available columns has been written to the partial results "
+                f"directory: {self.artifact_storage.partial_results_path}"
+            )
+
         self._current_batch_number += 1
         self._buffer: list[dict] = []
 
-        if on_complete:
-            on_complete(final_file_path)
-
         return final_file_path
 
     def finish(self) -> None:

data_designer/engine/dataset_builders/utils/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.engine.errors import DataDesignerError
 
 

data_designer/engine/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from pydantic import BaseModel, Field
 
 from data_designer.errors import DataDesignerError

data_designer/engine/model_provider.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from functools import cached_property
 
 from pydantic import BaseModel, field_validator, model_validator