data-designer 0.1.0__py3-none-any.whl

data_designer/engine/analysis/utils/judge_score_processing.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import defaultdict
+import logging
+from typing import Any, Optional, Union
+
+import pandas as pd
+
+from data_designer.config.analysis.column_profilers import JudgeScoreDistributions, JudgeScoreSample
+from data_designer.config.analysis.column_statistics import (
+    CategoricalDistribution,
+    ColumnDistributionType,
+    MissingValue,
+    NumericalDistribution,
+)
+from data_designer.config.column_configs import LLMJudgeColumnConfig
+
+logger = logging.getLogger(__name__)
+
+
+def extract_judge_score_distributions(
+    column_config: LLMJudgeColumnConfig, df: pd.DataFrame
+) -> Union[JudgeScoreDistributions, MissingValue]:
+    scores = defaultdict(list)
+    reasoning = defaultdict(list)
+
+    # Aggregate results as dicts of form {score_name: <result>}.
+    histograms = {}
+    distributions = {}
+    distribution_types = {}
+
+    for score in column_config.scores:
+        is_numerical = True
+        name = score.name.lower()
+        for results in df[column_config.name]:
+            try:
+                score = results[name].get("score", None)
+
+                if _can_be_converted_to_int(score):
+                    score = int(score)
+                else:
+                    score = str(score)
+                    is_numerical = False
+
+                scores[name].append(score)
+                reasoning[name].append(results[name].get("reasoning", "No reasoning provided"))
+            except Exception as e:
+                logger.warning(f"⚠️ Failed to extract judge score for '{name}': {e}")
+                return MissingValue.OUTPUT_FORMAT_ERROR
+
+        try:
+            series = pd.Series(scores[name], name=name)
+            cat_dist = CategoricalDistribution.from_series(series)
+
+            # For judge scores, build a categorical histogram, since numerical scores are integers.
+            histograms[name] = cat_dist.histogram
+
+            if is_numerical:
+                distribution_types[name] = ColumnDistributionType.NUMERICAL
+                distributions[name] = NumericalDistribution.from_series(series)
+            else:
+                distribution_types[name] = ColumnDistributionType.CATEGORICAL
+                distributions[name] = cat_dist
+
+        except Exception as e:
+            logger.warning(f"⚠️ Failed to calculate judge score distribution for '{name}': {e}")
+            distribution_types[name] = ColumnDistributionType.UNKNOWN
+            distributions[name] = MissingValue.CALCULATION_FAILED
+            histograms[name] = MissingValue.CALCULATION_FAILED
+
+    return JudgeScoreDistributions(
+        scores=dict(scores),
+        reasoning=dict(reasoning),
+        distribution_types=distribution_types,
+        distributions=distributions,
+        histograms=histograms,
+    )
+
+
+def sample_scores_and_reasoning(
+    scores: list[Union[int, str]],
+    reasoning: list[str],
+    num_samples: int,
+    random_seed: Optional[int] = None,
+) -> list[JudgeScoreSample]:
+    if len(scores) != len(reasoning):
+        raise ValueError("scores and reasoning must have the same length")
+
+    if len(scores) == 0:
+        raise ValueError("scores and reasoning must not be empty")
+
+    if num_samples <= 0:
+        raise ValueError("num_samples must be greater than 0")
+
+    df_samples = pd.DataFrame({"score": scores, "reasoning": reasoning})
+
+    if len(scores) <= num_samples:
+        return [JudgeScoreSample(score=score, reasoning=reasoning) for score, reasoning in zip(scores, reasoning)]
+
+    # Sample maintaining original proportions from each category (int or str)
+    # Calculate the frequency of each score category
+    score_category_counts = df_samples["score"].value_counts()
+
+    # If more categories than samples, pick one sample from each of the most frequent categories
+    if len(score_category_counts) >= num_samples:
+        top_categories = score_category_counts.head(num_samples).index
+        samples = pd.concat(
+            [df_samples[df_samples["score"] == cat].sample(n=1, random_state=random_seed) for cat in top_categories],
+            ignore_index=True,
+        )
+    else:
+        # Sample proportionally to maintain original category ratios
+        # Create weights based on the original frequency of each score
+        weights = df_samples["score"].map(score_category_counts)
+        samples = df_samples.sample(n=num_samples, weights=weights, random_state=random_seed)
+
+    return [
+        JudgeScoreSample(score=row["score"], reasoning=row["reasoning"]) for row in samples.to_dict(orient="records")
+    ]
+
+
+def _can_be_converted_to_int(value: Any) -> bool:
+    try:
+        int(value)
+        return True
+    except (ValueError, TypeError):
+        return False

data_designer/engine/column_generators/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

data_designer/engine/column_generators/generators/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

data_designer/engine/column_generators/generators/base.py

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from typing import overload
+
+import pandas as pd
+
+from data_designer.config.utils.type_helpers import StrEnum
+from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, DataT, TaskConfigT
+
+
+class GenerationStrategy(StrEnum):
+    CELL_BY_CELL = "cell_by_cell"
+    FULL_COLUMN = "full_column"
+
+
+class GeneratorMetadata(ConfigurableTaskMetadata):
+    generation_strategy: GenerationStrategy
+
+
+class ColumnGenerator(ConfigurableTask[TaskConfigT], ABC):
+    @property
+    def can_generate_from_scratch(self) -> bool:
+        return False
+
+    @property
+    def generation_strategy(self) -> GenerationStrategy:
+        return self.metadata().generation_strategy
+
+    @staticmethod
+    @abstractmethod
+    def metadata() -> GeneratorMetadata: ...
+
+    @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 log_pre_generation(self) -> None:
+        """A shared method to log info before the generator's `generate` method is called.
+
+        The idea is for dataset builders to call this method for all generators before calling their
+        `generate` method. This is to avoid logging the same information multiple times when running
+        generators in parallel.
+        """
+
+
+class FromScratchColumnGenerator(ColumnGenerator[TaskConfigT], ABC):
+    @property
+    def can_generate_from_scratch(self) -> bool:
+        return True
+
+    @abstractmethod
+    def generate_from_scratch(self, num_records: int) -> pd.DataFrame: ...

data_designer/engine/column_generators/generators/expression.py

@@ -0,0 +1,63 @@
+# 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 ExpressionColumnConfig
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    GeneratorMetadata,
+)
+from data_designer.engine.column_generators.utils.errors import ExpressionTemplateRenderError
+from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
+from data_designer.engine.processing.utils import deserialize_json_values
+
+logger = logging.getLogger(__name__)
+
+
+class ExpressionColumnGenerator(WithJinja2UserTemplateRendering, ColumnGenerator[ExpressionColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="expression_generator",
+            description="Generate a column from a jinja2 expression.",
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            required_resources=None,
+        )
+
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame:
+        logger.info(f"🧩 Generating column `{self.config.name}` from expression")
+
+        missing_columns = list(set(self.config.required_columns) - set(data.columns))
+        if len(missing_columns) > 0:
+            error_msg = (
+                f"There was an error preparing the Jinja2 expression template. "
+                f"The following columns {missing_columns} are missing!"
+            )
+            raise ExpressionTemplateRenderError(error_msg)
+
+        self.prepare_jinja2_template_renderer(self.config.expr, data.columns.to_list())
+        records = []
+        for record in data.to_dict(orient="records"):
+            record[self.config.name] = self._cast_type(self.render_template(deserialize_json_values(record)))
+            records.append(record)
+
+        return pd.DataFrame(records)
+
+    def _cast_type(self, value: str) -> str | float | int | bool:
+        if self.config.dtype == "str":
+            return value
+        elif self.config.dtype == "float":
+            return float(value)
+        elif self.config.dtype == "int":
+            return int(float(value))
+        elif self.config.dtype == "bool":
+            try:
+                return bool(int(float(value)))
+            except ValueError:
+                return bool(f"{value}".lower() == "true")
+        else:
+            raise ValueError(f"Invalid dtype: {self.config.dtype}")

data_designer/engine/column_generators/generators/llm_generators.py

@@ -0,0 +1,172 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import logging
+
+from data_designer.config.column_configs import (
+    LLMCodeColumnConfig,
+    LLMJudgeColumnConfig,
+    LLMStructuredColumnConfig,
+    LLMTextColumnConfig,
+)
+from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP
+from data_designer.config.models import InferenceParameters, ModelConfig
+from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    GeneratorMetadata,
+)
+from data_designer.engine.column_generators.utils.prompt_renderer import (
+    PromptType,
+    RecordBasedPromptRenderer,
+    create_response_recipe,
+)
+from data_designer.engine.models.facade import ModelFacade
+from data_designer.engine.models.recipes.base import ResponseRecipe
+from data_designer.engine.processing.utils import deserialize_json_values
+from data_designer.engine.resources.resource_provider import ResourceType
+
+DEFAULT_MAX_CONVERSATION_RESTARTS = 5
+DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS = 0
+
+
+logger = logging.getLogger(__name__)
+
+
+class WithLLMGeneration:
+    @functools.cached_property
+    def model(self) -> ModelFacade:
+        return self.resource_provider.model_registry.get_model(model_alias=self.config.model_alias)
+
+    @functools.cached_property
+    def model_config(self) -> ModelConfig:
+        return self.resource_provider.model_registry.get_model_config(model_alias=self.config.model_alias)
+
+    @functools.cached_property
+    def inference_parameters(self) -> InferenceParameters:
+        return self.model_config.inference_parameters
+
+    @functools.cached_property
+    def prompt_renderer(self) -> RecordBasedPromptRenderer:
+        return RecordBasedPromptRenderer(
+            response_recipe=self.response_recipe,
+            error_message_context={
+                "column_name": self.config.name,
+                "column_type": self.config.column_type,
+                "model_alias": self.config.model_alias,
+            },
+        )
+
+    @functools.cached_property
+    def response_recipe(self) -> ResponseRecipe:
+        return create_response_recipe(self.config, self.model_config)
+
+    @property
+    def max_conversation_correction_steps(self) -> int:
+        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
+
+    @property
+    def max_conversation_restarts(self) -> int:
+        return DEFAULT_MAX_CONVERSATION_RESTARTS
+
+    def generate(self, data: dict) -> dict:
+        deserialized_record = deserialize_json_values(data)
+
+        multi_modal_context = None
+        if self.config.multi_modal_context is not None and len(self.config.multi_modal_context) > 0:
+            multi_modal_context = [
+                context.get_context(deserialized_record) for context in self.config.multi_modal_context
+            ]
+
+        response, reasoning_trace = self.model.generate(
+            prompt=self.prompt_renderer.render(
+                record=deserialized_record,
+                prompt_template=self.config.prompt,
+                prompt_type=PromptType.USER_PROMPT,
+            ),
+            system_prompt=self.prompt_renderer.render(
+                record=deserialized_record,
+                prompt_template=self.config.system_prompt,
+                prompt_type=PromptType.SYSTEM_PROMPT,
+            ),
+            parser=self.response_recipe.parse,
+            multi_modal_context=multi_modal_context,
+            max_correction_steps=self.max_conversation_correction_steps,
+            max_conversation_restarts=self.max_conversation_restarts,
+            purpose=f"running generation for column '{self.config.name}'",
+            **self.inference_parameters.generate_kwargs,
+        )
+
+        data[self.config.name] = deserialize_json_values(self.response_recipe.serialize_output(response))
+
+        if reasoning_trace:
+            data[self.config.name + REASONING_TRACE_COLUMN_POSTFIX] = reasoning_trace
+
+        return data
+
+    def log_pre_generation(self) -> None:
+        emoji = COLUMN_TYPE_EMOJI_MAP[self.config.column_type]
+        logger.info(f"{emoji} Preparing {self.config.column_type} column generation")
+        logger.info(f"  |-- column name: {self.config.name!r}")
+        logger.info(f"  |-- model config:\n{self.model_config.model_dump_json(indent=4)}")
+        if self.model_config.provider is None:
+            logger.info(f"  |-- default model provider: {self._get_provider_name()!r}")
+
+    def _get_provider_name(self) -> str:
+        model_alias = self.model_config.alias
+        provider = self.resource_provider.model_registry.get_model_provider(model_alias=model_alias)
+        return provider.name
+
+
+class LLMTextCellGenerator(WithLLMGeneration, ColumnGenerator[LLMTextColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="llm_text_generator",
+            description="Generate a new dataset cell from a prompt template",
+            generation_strategy=GenerationStrategy.CELL_BY_CELL,
+            required_resources=[ResourceType.MODEL_REGISTRY],
+        )
+
+
+class LLMCodeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMCodeColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="llm_code_generator",
+            description="Generate a new dataset cell from a prompt template",
+            generation_strategy=GenerationStrategy.CELL_BY_CELL,
+            required_resources=[ResourceType.MODEL_REGISTRY],
+        )
+
+
+class LLMStructuredCellGenerator(WithLLMGeneration, ColumnGenerator[LLMStructuredColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="llm_structured_generator",
+            description="Generate a new dataset cell from a prompt template",
+            generation_strategy=GenerationStrategy.CELL_BY_CELL,
+            required_resources=[ResourceType.MODEL_REGISTRY],
+        )
+
+
+class LLMJudgeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMJudgeColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="llm_judge_generator",
+            description="Judge a new dataset cell based on a set of rubrics",
+            generation_strategy=GenerationStrategy.CELL_BY_CELL,
+            required_resources=[ResourceType.MODEL_REGISTRY],
+        )
+
+    @property
+    def max_conversation_correction_steps(self) -> int:
+        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
+
+    @property
+    def max_conversation_restarts(self) -> int:
+        return 2 * DEFAULT_MAX_CONVERSATION_RESTARTS

data_designer/engine/column_generators/generators/samplers.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import partial
+import logging
+import random
+from typing import Callable
+
+import pandas as pd
+
+from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
+from data_designer.engine.column_generators.generators.base import (
+    FromScratchColumnGenerator,
+    GenerationStrategy,
+    GeneratorMetadata,
+)
+from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
+from data_designer.engine.processing.utils import concat_datasets
+from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
+from data_designer.engine.resources.resource_provider import ResourceType
+from data_designer.engine.sampling_gen.data_sources.sources import SamplerType
+from data_designer.engine.sampling_gen.entities.person import load_person_data_sampler
+from data_designer.engine.sampling_gen.generator import DatasetGenerator as SamplingDatasetGenerator
+
+logger = logging.getLogger(__name__)
+
+
+class SamplerColumnGenerator(FromScratchColumnGenerator[SamplerMultiColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="sampler_column_generator",
+            description="Generate columns using sampling-based method.",
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            required_resources=[ResourceType.BLOB_STORAGE],
+        )
+
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame:
+        df_samplers = self.generate_from_scratch(len(data))
+        return concat_datasets([data, df_samplers])
+
+    def generate_from_scratch(self, num_records: int) -> pd.DataFrame:
+        sampling_generator = self._prepare_for_generation(num_records)
+        return sampling_generator.generate(num_records)
+
+    @property
+    def _needs_person_generator(self) -> bool:
+        columns = [c for c in self.config.columns if c.sampler_type == SamplerType.PERSON]
+        return any(c.params.locale in LOCALES_WITH_MANAGED_DATASETS for c in columns)
+
+    @property
+    def _person_generator_loader(self) -> Callable[[bool], ManagedDatasetGenerator]:
+        return partial(load_person_data_sampler, blob_storage=self.resource_provider.blob_storage)
+
+    def _create_sampling_dataset_generator(self) -> SamplingDatasetGenerator:
+        return SamplingDatasetGenerator(
+            sampler_columns=self.config,
+            person_generator_loader=(self._person_generator_loader if self._needs_person_generator else None),
+        )
+
+    def _log_person_generation_if_needed(self) -> None:
+        if self._needs_person_generator:
+            columns = [c for c in self.config.columns if c.sampler_type == SamplerType.PERSON]
+            emoji = random.choice(["🧑‍🎨", "🙋‍♂️", "🙋‍♀️", "🧑‍🚀", "👩‍🎤", "👨‍🍳", "👩‍🔬", "👨‍💻", "👩‍💼"])
+            log_msg = f"🎲 {emoji} Initializing person generation"
+            if any(c.params.with_synthetic_personas for c in columns):
+                log_msg += " ⚡️ with synthetic personas ⚡️"
+            logger.info(log_msg)
+
+    def _prepare_for_generation(self, num_records: int) -> SamplingDatasetGenerator:
+        logger.info(
+            f"🎲 Preparing samplers to generate {num_records} records across {len(self.config.columns)} columns"
+        )
+        self._log_person_generation_if_needed()
+        return self._create_sampling_dataset_generator()

data_designer/engine/column_generators/generators/seed_dataset.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import logging
+
+import duckdb
+import pandas as pd
+
+from data_designer.config.seed import IndexRange, PartitionBlock, SamplingStrategy
+from data_designer.engine.column_generators.generators.base import (
+    FromScratchColumnGenerator,
+    GenerationStrategy,
+    GeneratorMetadata,
+)
+from data_designer.engine.column_generators.utils.errors import SeedDatasetError
+from data_designer.engine.dataset_builders.multi_column_configs import SeedDatasetMultiColumnConfig
+from data_designer.engine.processing.utils import concat_datasets
+from data_designer.engine.resources.resource_provider import ResourceType
+
+MAX_ZERO_RECORD_RESPONSE_FACTOR = 2
+
+logger = logging.getLogger(__name__)
+
+
+class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColumnConfig]):
+    @staticmethod
+    def metadata() -> GeneratorMetadata:
+        return GeneratorMetadata(
+            name="seed_dataset_column_generator",
+            description="Sample columns from a seed dataset.",
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            required_resources=[ResourceType.DATASTORE],
+        )
+
+    @property
+    def num_records_sampled(self) -> int:
+        return self._num_records_sampled
+
+    @functools.cached_property
+    def duckdb_conn(self) -> duckdb.DuckDBPyConnection:
+        return self.resource_provider.datastore.create_duckdb_connection()
+
+    def generate(self, dataset: pd.DataFrame) -> pd.DataFrame:
+        return concat_datasets([self.generate_from_scratch(len(dataset)), dataset])
+
+    def generate_from_scratch(self, num_records: int) -> pd.DataFrame:
+        if num_records <= 0:
+            raise ValueError("🛑 `num_records` must be positive.")
+
+        if self._batch_reader is None:
+            self._reset_batch_reader(num_records)
+
+        return self._sample_records(num_records)
+
+    def _initialize(self) -> None:
+        self._num_records_sampled = 0
+        self._batch_reader = None
+        self._df_remaining = None
+        self._dataset_uri = self.resource_provider.datastore.get_dataset_uri(self.config.dataset)
+        self._seed_dataset_size = self.duckdb_conn.execute(f"SELECT COUNT(*) FROM '{self._dataset_uri}'").fetchone()[0]
+        self._index_range = self._resolve_index_range()
+
+    def _validate_selection_strategy(self) -> None:
+        err_msg = None
+        if self.config.selection_strategy is not None:
+            if (
+                isinstance(self.config.selection_strategy, IndexRange)
+                and self.config.selection_strategy.end >= self._seed_dataset_size
+            ):
+                err_msg = f"Selection strategy 'end' index {self.config.selection_strategy.end} is out of bounds for dataset size {self._seed_dataset_size}"
+            elif (
+                isinstance(self.config.selection_strategy, PartitionBlock)
+                and self.config.selection_strategy.num_partitions > self._seed_dataset_size
+            ):
+                err_msg = f"Selection strategy 'num_partitions' {self.config.selection_strategy.num_partitions} is out of bounds for dataset size {self._seed_dataset_size}"
+            if err_msg is not None:
+                raise SeedDatasetError(err_msg)
+
+    def _resolve_index_range(self) -> IndexRange | None:
+        self._validate_selection_strategy()
+        index_range = None
+        if self.config.selection_strategy is not None:
+            if isinstance(self.config.selection_strategy, IndexRange):
+                index_range = self.config.selection_strategy
+            elif isinstance(self.config.selection_strategy, PartitionBlock):
+                index_range = self.config.selection_strategy.to_index_range(self._seed_dataset_size)
+        return index_range
+
+    def _reset_batch_reader(self, num_records: int) -> None:
+        shuffle = self.config.sampling_strategy == SamplingStrategy.SHUFFLE
+        shuffle_query = " ORDER BY RANDOM()" if shuffle else ""
+
+        if self._index_range is not None:
+            # Use LIMIT and OFFSET for efficient index range filtering
+            # IndexRange uses 0-based indexing [start, end] inclusive
+            # OFFSET skips the first 'start' rows (0-based)
+            # LIMIT takes 'end - start + 1' rows to include both start and end (inclusive)
+            offset_value = self._index_range.start
+            limit_value = self._index_range.end - self._index_range.start + 1
+            read_query = f"""
+                SELECT * FROM '{self._dataset_uri}'
+                LIMIT {limit_value} OFFSET {offset_value}
+            """
+
+            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)
+
+    def _sample_records(self, num_records: int) -> pd.DataFrame:
+        logger.info(f"🌱 Sampling {num_records} records from seed dataset")
+        logger.info(f"  |-- seed dataset size: {self._seed_dataset_size} records")
+        logger.info(f"  |-- sampling strategy: {self.config.sampling_strategy}")
+        if self._index_range is not None:
+            if isinstance(self.config.selection_strategy, IndexRange):
+                logger.info(f"  |-- selection: rows [{self._index_range.start} to {self._index_range.end}] inclusive")
+            else:
+                logger.info(
+                    f"  |-- selection: partition {self.config.selection_strategy.index + 1} of {self.config.selection_strategy.num_partitions}"
+                )
+            logger.info(f"  |-- seed dataset size after selection: {self._index_range.size} records")
+        df_batch = pd.DataFrame()
+        df_sample = pd.DataFrame() if self._df_remaining is None else self._df_remaining
+        num_zero_record_responses = 0
+
+        while len(df_sample) < num_records:
+            try:
+                df_batch = self._batch_reader.read_next_batch().to_pandas()
+                df_sample = pd.concat([df_sample, df_batch], ignore_index=True)
+            except StopIteration:
+                self._reset_batch_reader(num_records)
+
+            if len(df_batch) == 0:
+                num_zero_record_responses += 1
+                if num_zero_record_responses > MAX_ZERO_RECORD_RESPONSE_FACTOR * num_records:
+                    raise RuntimeError(
+                        "🛑 Something went wrong while reading from the datastore. "
+                        "Please check your connection and try again. "
+                        "If the issue persists, please contact support."
+                    )
+
+        self._df_remaining = None
+        if len(df_sample) > num_records:
+            self._df_remaining = df_sample.iloc[num_records:].reset_index(drop=True)
+        df_sample = df_sample.iloc[:num_records]
+        self._num_records_sampled += len(df_sample)
+
+        return df_sample