data-designer-engine 0.5.2__tar.gz → 0.5.3__tar.gz

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/.gitignore

@@ -99,8 +99,14 @@ NOTEPAD.md
 # Build-time copy of README for data-designer package (copied from top-level during build)
 packages/data-designer/README.md
 
+# Notebook build cache
+.notebook-cache/
+
 # Cerebro knowledge base
 .cerebro/
 .cursor/rules/cerebro.mdc
 .cursor/mcp.json
 .claude/rules/cerebro.md
+
+# Claude worktrees
+.claude/worktrees/

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer-engine
-Version: 0.5.2
+Version: 0.5.3
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta
@@ -14,8 +14,9 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Requires-Dist: anyascii<1,>=0.3.3
-Requires-Dist: data-designer-config==0.5.2
-Requires-Dist: duckdb<2,>=1.1.3
+Requires-Dist: chardet<6,>=3.0.2
+Requires-Dist: data-designer-config==0.5.3
+Requires-Dist: duckdb<2,>=1.5.0
 Requires-Dist: faker<21,>=20.1.0
 Requires-Dist: httpx-retries<1,>=0.4.2
 Requires-Dist: httpx<1,>=0.27.2

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/pyproject.toml

@@ -33,8 +33,9 @@ bump = true
 [tool.hatch.metadata.hooks.uv-dynamic-versioning]
 dependencies = [
     "anyascii>=0.3.3,<1",
+    "chardet>=3.0.2,<6",  # Pulled in by sqlfluff; pin <6 to avoid RequestsDependencyWarning from requests<2.33
     "data-designer-config=={{ version }}",
-    "duckdb>=1.1.3,<2",
+    "duckdb>=1.5.0,<2",
     "faker>=20.1.0,<21",
     "httpx>=0.27.2,<1",
     "httpx-retries>=0.4.2,<1",

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/src/data_designer/engine/column_generators/generators/base.py

@@ -4,15 +4,20 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import functools
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, overload
+from typing import TYPE_CHECKING, Any, Coroutine, TypeVar, overload
 
 from data_designer.config.column_configs import GenerationStrategy
 from data_designer.engine.configurable_task import ConfigurableTask, DataT, TaskConfigT
 from data_designer.logging import LOG_DOUBLE_INDENT, LOG_INDENT
 
+_T = TypeVar("_T")
+
+_SYNC_BRIDGE_TIMEOUT = 300
+
 if TYPE_CHECKING:
     import pandas as pd
 
@@ -23,33 +28,84 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _run_coroutine_sync(coro: Coroutine[Any, Any, _T]) -> _T:
+    """Run an async coroutine from sync context.
+
+    - No running event loop → ``asyncio.run(coro)``
+    - Running event loop (e.g. notebook/service) → run in a background thread
+    """
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(coro)
+    pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+    future = pool.submit(asyncio.run, coro)
+    timed_out = False
+    try:
+        result = future.result(timeout=_SYNC_BRIDGE_TIMEOUT)
+    except concurrent.futures.TimeoutError as exc:
+        timed_out = True
+        logger.warning(f"⚠️ Sync bridge timed out after {_SYNC_BRIDGE_TIMEOUT}s; background thread still running")
+        raise TimeoutError(f"_run_coroutine_sync timed out after {_SYNC_BRIDGE_TIMEOUT}s") from exc
+    finally:
+        pool.shutdown(wait=not timed_out, cancel_futures=timed_out)
+    return result
+
+
 class ColumnGenerator(ConfigurableTask[TaskConfigT], ABC):
     @property
     def can_generate_from_scratch(self) -> bool:
         return False
 
+    @property
+    def is_order_dependent(self) -> bool:
+        """Whether this generator's output depends on prior row-group calls.
+
+        Example: SeedDatasetColumnGenerator tracks its position in the seed
+        dataset, so row group N must complete before N+1 starts.
+        """
+        return False
+
+    def _is_overridden(self, method_name: str) -> bool:
+        """Check if a subclass has overridden a base ColumnGenerator method."""
+        return getattr(type(self), method_name) is not getattr(ColumnGenerator, method_name)
+
     @staticmethod
     @abstractmethod
     def get_generation_strategy() -> GenerationStrategy: ...
 
     @overload
-    @abstractmethod
     def generate(self, data: dict) -> dict: ...
 
     @overload
-    @abstractmethod
     def generate(self, data: pd.DataFrame) -> pd.DataFrame: ...
 
-    @abstractmethod
-    def generate(self, data: DataT) -> DataT: ...
+    def generate(self, data: DataT) -> DataT:
+        """Sync generate — overridden by most concrete generators.
+
+        Default bridges to ``agenerate()`` for async-first subclasses that only
+        implement ``agenerate()``. Raises ``NotImplementedError`` if neither
+        ``generate()`` nor ``agenerate()`` is overridden.
+        """
+        if not self._is_overridden("agenerate"):
+            raise NotImplementedError(f"{type(self).__name__} must implement either generate() or agenerate()")
+        return _run_coroutine_sync(self.agenerate(data))
 
-    async def agenerate(self, data: dict) -> dict:
-        """Async fallback — delegates to sync generate via thread pool.
+    @overload
+    async def agenerate(self, data: dict) -> dict: ...
+
+    @overload
+    async def agenerate(self, data: pd.DataFrame) -> pd.DataFrame: ...
+
+    async def agenerate(self, data: DataT) -> DataT:
+        """Async generate — delegates to sync ``generate()`` via thread pool.
 
         Subclasses with native async support (e.g. ColumnGeneratorWithModelChatCompletion)
         should override this with a direct async implementation.
         """
-        return await asyncio.to_thread(self.generate, data)
+        if not self._is_overridden("generate"):
+            raise NotImplementedError(f"{type(self).__name__} must implement either generate() or agenerate()")
+        return await asyncio.to_thread(self.generate, data.copy())
 
     def log_pre_generation(self) -> None:
         """A shared method to log info before the generator's `generate` method is called.
@@ -68,6 +124,10 @@ class FromScratchColumnGenerator(ColumnGenerator[TaskConfigT], ABC):
     @abstractmethod
     def generate_from_scratch(self, num_records: int) -> pd.DataFrame: ...
 
+    async def agenerate_from_scratch(self, num_records: int) -> pd.DataFrame:
+        """Async wrapper — wraps sync ``generate_from_scratch()`` in a thread."""
+        return await asyncio.to_thread(self.generate_from_scratch, num_records)
+
 
 class ColumnGeneratorWithModelRegistry(ColumnGenerator[TaskConfigT], ABC):
     @property

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/src/data_designer/engine/column_generators/generators/custom.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import inspect
 import logging
 from typing import TYPE_CHECKING, Any
@@ -65,12 +66,57 @@ class CustomColumnGenerator(ColumnGenerator[CustomColumnConfig]):
 
         return self._generate(data, is_dataframe)
 
+    async def agenerate(self, data: dict | pd.DataFrame) -> dict | pd.DataFrame | list[dict]:
+        """Async generate — branches on strategy and detects coroutine functions."""
+        is_full_column = self.config.generation_strategy == GenerationStrategy.FULL_COLUMN
+        if is_full_column:
+            return await asyncio.to_thread(self.generate, data.copy())
+        # The @custom_column_generator decorator wraps the user function in a sync
+        # wrapper, so we must unwrap to detect async functions.
+        fn_unwrapped = inspect.unwrap(self.config.generator_function)
+        if asyncio.iscoroutinefunction(fn_unwrapped):
+            missing = set(self.config.required_columns) - set(data.keys())
+            if missing:
+                raise CustomColumnGenerationError(
+                    f"Missing required columns for custom generator '{self.config.name}': {sorted(missing)}"
+                )
+            keys_before = set(data.keys())
+
+            try:
+                result = await self._ainvoke_generator_function(data)
+            except CustomColumnGenerationError:
+                raise
+            except Exception as e:
+                logger.warning(
+                    f"⚠️ Custom generator function {self.config.generator_function.__name__!r} "
+                    f"failed for column '{self.config.name}'. This record will be skipped.\n{e}"
+                )
+                raise CustomColumnGenerationError(
+                    f"Custom generator function failed for column '{self.config.name}': {e}"
+                ) from e
+
+            return self._postprocess_result(result, is_dataframe=False, keys_before=keys_before)
+        return await asyncio.to_thread(self.generate, data)
+
+    async def _ainvoke_generator_function(self, data: dict) -> dict | pd.DataFrame:
+        """Invoke an async user generator function with appropriate arguments.
+
+        The @custom_column_generator decorator's sync wrapper returns a coroutine
+        when the original function is async, so we await the wrapper's return value.
+        """
+        params = self._get_validated_params()
+        fn = self.config.generator_function
+        if len(params) == 1:
+            return await fn(data)
+        elif len(params) == 2:
+            return await fn(data, self.config.generator_params)
+        else:
+            models = self._build_models_dict()
+            return await fn(data, self.config.generator_params, models)
+
     def _generate(self, data: dict | pd.DataFrame, is_dataframe: bool) -> dict | pd.DataFrame | list[dict]:
         """Unified generation logic for both strategies."""
-        # Get columns/keys using unified accessor
         get_keys = (lambda d: set(d.columns)) if is_dataframe else (lambda d: set(d.keys()))
-        expected_type = lazy.pd.DataFrame if is_dataframe else dict
-        type_name = "DataFrame" if is_dataframe else "dict"
 
         # Check required columns
         missing = set(self.config.required_columns) - get_keys(data)
@@ -96,6 +142,15 @@ class CustomColumnGenerator(ColumnGenerator[CustomColumnConfig]):
                 f"Custom generator function failed for column '{self.config.name}': {e}"
             ) from e
 
+        return self._postprocess_result(result, is_dataframe, keys_before)
+
+    def _postprocess_result(
+        self,
+        result: dict | pd.DataFrame | list[dict],
+        is_dataframe: bool,
+        keys_before: set[str],
+    ) -> dict | pd.DataFrame | list[dict]:
+        """Validate type and output columns of a generation result."""
         # Cell-by-cell with allow_resize: accept dict or list[dict]
         if not is_dataframe and self.config.allow_resize:
             if isinstance(result, dict):
@@ -113,6 +168,8 @@ class CustomColumnGenerator(ColumnGenerator[CustomColumnConfig]):
             )
 
         # Validate return type for non-resize paths
+        expected_type = lazy.pd.DataFrame if is_dataframe else dict
+        type_name = "DataFrame" if is_dataframe else "dict"
         if not isinstance(result, expected_type):
             raise CustomColumnGenerationError(
                 f"Custom generator for column '{self.config.name}' must return a {type_name}, "

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/src/data_designer/engine/column_generators/generators/embedding.py

@@ -27,9 +27,19 @@ class EmbeddingCellGenerator(ColumnGeneratorWithModel[EmbeddingColumnConfig]):
     def get_generation_strategy() -> GenerationStrategy:
         return GenerationStrategy.CELL_BY_CELL
 
-    def generate(self, data: dict) -> dict:
+    def _prepare_embedding_inputs(self, data: dict) -> list[str]:
         deserialized_record = deserialize_json_values(data)
-        input_texts = parse_list_string(deserialized_record[self.config.target_column])
+        return parse_list_string(deserialized_record[self.config.target_column])
+
+    def generate(self, data: dict) -> dict:
+        input_texts = self._prepare_embedding_inputs(data)
         embeddings = self.model.generate_text_embeddings(input_texts=input_texts)
         data[self.config.name] = EmbeddingGenerationResult(embeddings=embeddings).model_dump(mode="json")
         return data
+
+    async def agenerate(self, data: dict) -> dict:
+        """Native async generate using model.agenerate_text_embeddings."""
+        input_texts = self._prepare_embedding_inputs(data)
+        embeddings = await self.model.agenerate_text_embeddings(input_texts=input_texts)
+        data[self.config.name] = EmbeddingGenerationResult(embeddings=embeddings).model_dump(mode="json")
+        return data

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/src/data_designer/engine/column_generators/generators/image.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from typing import TYPE_CHECKING
 
 from data_designer.config.column_configs import ImageColumnConfig
@@ -31,46 +32,42 @@ class ImageCellGenerator(WithJinja2UserTemplateRendering, ColumnGeneratorWithMod
     def get_generation_strategy() -> GenerationStrategy:
         return GenerationStrategy.CELL_BY_CELL
 
-    def generate(self, data: dict) -> dict:
-        """Generate image(s) and optionally save to disk.
-
-        Args:
-            data: Record data
-
-        Returns:
-            Record with image path(s) (create mode) or base64 data (preview mode) added
-        """
+    def _prepare_image_inputs(self, data: dict) -> tuple[str, list[dict] | None]:
+        """Validate inputs and render prompt for image generation."""
         deserialized_record = deserialize_json_values(data)
-
-        # Validate required columns
         missing_columns = list(set(self.config.required_columns) - set(data.keys()))
         if len(missing_columns) > 0:
-            error_msg = (
+            raise ValueError(
                 f"There was an error preparing the Jinja2 expression template. "
                 f"The following columns {missing_columns} are missing!"
             )
-            raise ValueError(error_msg)
-
-        # Render prompt template
         self.prepare_jinja2_template_renderer(self.config.prompt, list(deserialized_record.keys()))
         prompt = self.render_template(deserialized_record)
-
-        # Validate prompt is non-empty
         if not prompt or not prompt.strip():
             raise ValueError(f"Rendered prompt for column {self.config.name!r} is empty")
-
-        # Process multi-modal context if provided
         multi_modal_context = self._build_multi_modal_context(deserialized_record)
+        return prompt, multi_modal_context
 
-        # Generate images (returns list of base64 strings)
+    def generate(self, data: dict) -> dict:
+        """Generate image(s) and optionally save to disk."""
+        prompt, multi_modal_context = self._prepare_image_inputs(data)
         base64_images = self.model.generate_image(prompt=prompt, multi_modal_context=multi_modal_context)
-
-        # Store via media storage (mode determines disk vs dataframe storage)
-        # Use column name as subfolder to organize images
         results = [
             self.media_storage.save_base64_image(base64_image, subfolder_name=self.config.name)
             for base64_image in base64_images
         ]
         data[self.config.name] = results
+        return data
 
+    async def agenerate(self, data: dict) -> dict:
+        """Native async generate using model.agenerate_image."""
+        prompt, multi_modal_context = self._prepare_image_inputs(data)
+        base64_images = await self.model.agenerate_image(prompt=prompt, multi_modal_context=multi_modal_context)
+        results = await asyncio.to_thread(
+            lambda: [
+                self.media_storage.save_base64_image(base64_image, subfolder_name=self.config.name)
+                for base64_image in base64_images
+            ]
+        )
+        data[self.config.name] = results
         return data

{data_designer_engine-0.5.2 → data_designer_engine-0.5.3}/src/data_designer/engine/column_generators/generators/seed_dataset.py

@@ -29,6 +29,10 @@ class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColu
     def get_generation_strategy() -> GenerationStrategy:
         return GenerationStrategy.FULL_COLUMN
 
+    @property
+    def is_order_dependent(self) -> bool:
+        return True
+
     @property
     def num_records_sampled(self) -> int:
         return self._num_records_sampled
@@ -102,7 +106,7 @@ class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColu
             read_query = f"SELECT * FROM ({read_query}){shuffle_query}"
         else:
             read_query = f"SELECT * FROM '{self._dataset_uri}'{shuffle_query}"
-        self._batch_reader = self.duckdb_conn.query(read_query).record_batch(batch_size=num_records)
+        self._batch_reader = self.duckdb_conn.query(read_query).to_arrow_reader(batch_size=num_records)
 
     def _sample_records(self, num_records: int) -> pd.DataFrame:
         logger.info(f"🌱 Sampling {num_records} records from seed dataset")

data_designer_engine-0.5.3/src/data_designer/engine/dataset_builders/utils/completion_tracker.py

@@ -0,0 +1,230 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from data_designer.config.column_configs import GenerationStrategy
+from data_designer.engine.dataset_builders.utils.task_model import SliceRef, Task
+
+if TYPE_CHECKING:
+    from data_designer.engine.dataset_builders.utils.execution_graph import ExecutionGraph
+
+
+class CompletionTracker:
+    """Tracks which cells (column, row_group, row_index) are done.
+
+    Row indices are local to their row group (0-based).
+
+    Use ``with_graph`` to create a frontier-enabled tracker where
+    ``get_ready_tasks`` returns in O(frontier) instead of scanning all
+    columns x rows x row groups.
+    """
+
+    def __init__(self) -> None:
+        # row_group → column → set of completed local row indices
+        self._completed: dict[int, dict[str, set[int]]] = defaultdict(lambda: defaultdict(set))
+        # row_group → set of dropped row indices
+        self._dropped: dict[int, set[int]] = defaultdict(set)
+
+        self._graph: ExecutionGraph | None = None
+        self._row_group_sizes: dict[int, int] = {}
+        self._batch_complete: dict[int, set[str]] = defaultdict(set)
+        self._frontier: set[Task] = set()
+
+    @classmethod
+    def with_graph(cls, graph: ExecutionGraph, row_groups: list[tuple[int, int]]) -> CompletionTracker:
+        """Create a frontier-enabled tracker backed by an execution graph."""
+        tracker = cls()
+        tracker._graph = graph
+        tracker._row_group_sizes = {rg_id: size for rg_id, size in row_groups}
+        tracker._seed_frontier()
+        return tracker
+
+    def mark_cell_complete(self, column: str, row_group: int, row_index: int) -> None:
+        self._validate_row_group(row_group)
+        self._validate_strategy(column, GenerationStrategy.CELL_BY_CELL, "mark_cell_complete")
+        self._completed[row_group][column].add(row_index)
+        if self._graph is not None:
+            self._frontier.discard(Task(column=column, row_group=row_group, row_index=row_index, task_type="cell"))
+            self._enqueue_downstream(column, row_group, row_index=row_index)
+
+    def mark_row_range_complete(self, column: str, row_group: int, row_group_size: int) -> None:
+        expected = self._validate_row_group(row_group)
+        self._validate_strategy(column, GenerationStrategy.FULL_COLUMN, "mark_row_range_complete")
+        if expected is not None and row_group_size != expected:
+            raise ValueError(f"Row-group size mismatch for rg={row_group}: got {row_group_size}, expected {expected}")
+        self._completed[row_group][column] = set(range(row_group_size))
+        self._batch_complete[row_group].add(column)
+        if self._graph is not None:
+            self._frontier.discard(Task(column=column, row_group=row_group, row_index=None, task_type="batch"))
+            self._enqueue_downstream(column, row_group, row_index=None)
+
+    def is_complete(self, ref: SliceRef) -> bool:
+        return ref.row_index in self._completed.get(ref.row_group, {}).get(ref.column, set())
+
+    def is_all_complete(self, cells: list[SliceRef]) -> bool:
+        """Check whether all the given cells are done.
+
+        A ``row_index`` of ``None`` means the entire batch for that column must
+        have been completed via ``mark_row_range_complete``.
+        """
+        for ref in cells:
+            if ref.row_index is None:
+                if ref.column not in self._batch_complete.get(ref.row_group, set()):
+                    return False
+            elif not self.is_complete(ref):
+                return False
+        return True
+
+    def drop_row(self, row_group: int, row_index: int) -> None:
+        self._validate_row_group(row_group)
+        self._dropped[row_group].add(row_index)
+        if self._graph is not None:
+            # Remove cell tasks for this row from the frontier
+            for col in self._graph.columns:
+                self._frontier.discard(Task(column=col, row_group=row_group, row_index=row_index, task_type="cell"))
+            # Dropping a row may unblock batch downstream tasks
+            self._reevaluate_batch_tasks(row_group)
+
+    def is_dropped(self, row_group: int, row_index: int) -> bool:
+        return row_index in self._dropped.get(row_group, set())
+
+    def is_row_group_complete(
+        self,
+        row_group: int,
+        row_group_size: int,
+        all_columns: list[str],
+    ) -> bool:
+        """All non-dropped rows have all columns done."""
+        dropped = self._dropped.get(row_group, set())
+        completed = self._completed.get(row_group, {})
+        for ri in range(row_group_size):
+            if ri in dropped:
+                continue
+            for col in all_columns:
+                if ri not in completed.get(col, set()):
+                    return False
+        return True
+
+    def get_ready_tasks(self, dispatched: set[Task]) -> list[Task]:
+        """Return all currently dispatchable tasks from the frontier.
+
+        Excludes already-dispatched/in-flight tasks.
+        """
+        return [t for t in self._frontier if t not in dispatched]
+
+    def _seed_frontier(self) -> None:
+        """Populate the frontier with root tasks (columns with no upstream deps)."""
+        if self._graph is None:
+            raise RuntimeError("This method requires a graph to be set.")
+        for col in self._graph.get_root_columns():
+            strategy = self._graph.get_strategy(col)
+            for rg_id, rg_size in self._row_group_sizes.items():
+                if strategy == GenerationStrategy.CELL_BY_CELL:
+                    for ri in range(rg_size):
+                        self._frontier.add(Task(column=col, row_group=rg_id, row_index=ri, task_type="cell"))
+                else:
+                    self._frontier.add(Task(column=col, row_group=rg_id, row_index=None, task_type="batch"))
+
+    def _enqueue_downstream(self, column: str, row_group: int, row_index: int | None) -> None:
+        """Add newly-ready downstream tasks to the frontier."""
+        if self._graph is None:
+            raise RuntimeError("This method requires a graph to be set.")
+        rg_completed = self._completed.get(row_group, {})
+        rg_dropped = self._dropped.get(row_group, set())
+        rg_batch_complete = self._batch_complete.get(row_group, set())
+        rg_size = self._row_group_sizes[row_group]
+
+        for down in self._graph.get_downstream_columns(column):
+            batch_ups, cell_ups = self._graph.split_upstream_by_strategy(down)
+
+            if any(up not in rg_batch_complete for up in batch_ups):
+                continue
+
+            down_strategy = self._graph.get_strategy(down)
+
+            if down_strategy == GenerationStrategy.CELL_BY_CELL:
+                cell_up_completed = [rg_completed.get(up, set()) for up in cell_ups]
+                if row_index is not None:
+                    # Cell completion: only check the same row
+                    down_completed = rg_completed.get(down, set())
+                    if (
+                        row_index not in rg_dropped
+                        and row_index not in down_completed
+                        and all(row_index in s for s in cell_up_completed)
+                    ):
+                        task = Task(column=down, row_group=row_group, row_index=row_index, task_type="cell")
+                        self._frontier.add(task)
+                else:
+                    # Batch completion: check all non-dropped, non-complete rows
+                    down_completed = rg_completed.get(down, set())
+                    for ri in range(rg_size):
+                        if ri in rg_dropped or ri in down_completed:
+                            continue
+                        if all(ri in s for s in cell_up_completed):
+                            task = Task(column=down, row_group=row_group, row_index=ri, task_type="cell")
+                            self._frontier.add(task)
+            else:
+                # FULL_COLUMN downstream: ready when all cell upstreams are fully complete
+                if down not in rg_batch_complete and self._are_cell_ups_complete(
+                    cell_ups, rg_completed, rg_size, rg_dropped
+                ):
+                    task = Task(column=down, row_group=row_group, row_index=None, task_type="batch")
+                    self._frontier.add(task)
+
+    def _reevaluate_batch_tasks(self, row_group: int) -> None:
+        """Check if any batch tasks became ready after a row was dropped."""
+        if self._graph is None:
+            raise RuntimeError("This method requires a graph to be set.")
+        rg_completed = self._completed.get(row_group, {})
+        rg_dropped = self._dropped.get(row_group, set())
+        rg_batch_complete = self._batch_complete.get(row_group, set())
+        rg_size = self._row_group_sizes[row_group]
+
+        for col in self._graph.get_topological_order():
+            if self._graph.get_strategy(col) != GenerationStrategy.FULL_COLUMN:
+                continue
+            if col in rg_batch_complete:
+                continue
+            batch_ups, cell_ups = self._graph.split_upstream_by_strategy(col)
+            if any(up not in rg_batch_complete for up in batch_ups):
+                continue
+            if self._are_cell_ups_complete(cell_ups, rg_completed, rg_size, rg_dropped):
+                task = Task(column=col, row_group=row_group, row_index=None, task_type="batch")
+                self._frontier.add(task)
+
+    def _are_cell_ups_complete(
+        self,
+        cell_ups: list[str],
+        rg_completed: dict[str, set[int]],
+        rg_size: int,
+        rg_dropped: set[int],
+    ) -> bool:
+        """Check all non-dropped rows are complete for each cell-by-cell upstream column."""
+        for up in cell_ups:
+            up_completed = rg_completed.get(up, set())
+            for ri in range(rg_size):
+                if ri not in rg_dropped and ri not in up_completed:
+                    return False
+        return True
+
+    def _validate_strategy(self, column: str, expected: GenerationStrategy, method: str) -> None:
+        """Validate that *column* matches the expected strategy in graph-enabled mode."""
+        if self._graph is None:
+            return
+        actual = self._graph.get_strategy(column)
+        if actual != expected:
+            raise ValueError(f"{method}() requires {expected.value} strategy, but column '{column}' has {actual.value}")
+
+    def _validate_row_group(self, row_group: int) -> int | None:
+        """Validate row-group id in graph-enabled mode and return its expected size."""
+        if self._graph is None:
+            return None
+        expected = self._row_group_sizes.get(row_group)
+        if expected is None:
+            known = sorted(self._row_group_sizes)
+            raise ValueError(f"Unknown row_group {row_group}. Known row_groups: {known}")
+        return expected