⚡️ Speed up method EvaluationDataset._to_pyfunc_dataset by 15%

[codeflash-ai]

📄 15% (0.15x) speedup for EvaluationDataset._to_pyfunc_dataset in mlflow/genai/datasets/evaluation_dataset.py

⏱️ Runtime : 16.7 milliseconds → 14.5 milliseconds (best of 20 runs)

📝 Explanation and details

The optimized code achieves a 15% speedup through two key optimizations that reduce expensive repeated operations:

1. Attribute Lookup Caching
The original code calls self.name and self.digest every time to_evaluation_dataset() is invoked. Based on the read-only dependency code, these trigger the __getattr__ method which performs dynamic attribute delegation - checking both _mlflow_dataset and _databricks_dataset with hasattr() and getattr() calls. The optimization pre-fetches and caches these values as _cached_name and _cached_digest during initialization, eliminating ~8.3% of runtime spent on attribute access (7.82ms → 4.4ns in the profiler).

2. Import Statement Caching
The original code imports LegacyEvaluationDataset on every method call. While the import itself is fast, the optimization caches the imported class as self._legacy_eval_cls after the first use, avoiding repeated import overhead. This is particularly beneficial when the method is called multiple times.

Performance Impact by Test Case
The optimizations show the greatest benefit for:

• Edge cases with None values: 167% faster (empty dataframes, missing attributes)
• Large datasets: 3-6% faster for substantial DataFrames
• Repeated calls: 2-3% faster on subsequent invocations due to import caching

The optimizations are most effective when to_evaluation_dataset() is called frequently (common in evaluation loops) or when the underlying dataset's attribute access is expensive due to the delegation pattern. The caching approach maintains full behavioral compatibility while eliminating redundant computations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 41 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pandas as pd
# imports
import pytest
from mlflow.genai.datasets.evaluation_dataset import EvaluationDataset


class LegacyEvaluationDataset:
    """
    Dummy replacement for mlflow.data.evaluation_dataset.EvaluationDataset
    """
    def __init__(self, data, path=None, feature_names=None, name=None, digest=None):
        self.data = data
        self.path = path
        self.feature_names = feature_names
        self.name = name
        self.digest = digest

    def __eq__(self, other):
        # For testing, equality is based on all attributes
        if not isinstance(other, LegacyEvaluationDataset):
            return False
        return (
            self.data.equals(other.data)
            and self.path == other.path
            and self.feature_names == other.feature_names
            and self.name == other.name
            and self.digest == other.digest
        )

class _EntityEvaluationDataset:
    """
    Dummy replacement for mlflow.entities.evaluation_dataset.EvaluationDataset
    """
    def __init__(self, df, name=None, digest=None):
        self._df = df
        self.name = name
        self.digest = digest

    def to_df(self):
        return self._df

    def __eq__(self, other):
        if not isinstance(other, _EntityEvaluationDataset):
            return False
        return (
            self._df.equals(other._df)
            and self.name == other.name
            and self.digest == other.digest
        )

class DatabricksDataset:
    """
    Dummy replacement for databricks managed dataset.
    """
    def __init__(self, df, name=None, digest=None):
        self._df = df
        self.name = name
        self.digest = digest

    def to_df(self):
        return self._df

    def __eq__(self, other):
        if not isinstance(other, DatabricksDataset):
            return False
        return (
            self._df.equals(other._df)
            and self.name == other.name
            and self.digest == other.digest
        )

# ---- UNIT TESTS ----

# 1. BASIC TEST CASES

def test_basic_mlflow_dataset_conversion():
    # Test conversion from MLflow entity dataset
    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
    entity = _EntityEvaluationDataset(df, name="mlflow_ds", digest="abc123")
    eval_ds = EvaluationDataset(entity)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 574μs -> 575μs (0.075% slower)

def test_basic_databricks_dataset_conversion():
    # Test conversion from Databricks managed dataset
    df = pd.DataFrame({"x": [5, 6], "y": [7, 8]})
    dbr_ds = DatabricksDataset(df, name="dbr_ds", digest="xyz789")
    eval_ds = EvaluationDataset(dbr_ds)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 565μs -> 557μs (1.43% faster)

def test_to_evaluation_dataset_with_args():
    # Test passing path and feature_names
    df = pd.DataFrame({"f1": [1, 2], "f2": [3, 4]})
    entity = _EntityEvaluationDataset(df, name="test_ds", digest="digestval")
    eval_ds = EvaluationDataset(entity)
    pyfunc_ds = eval_ds.to_evaluation_dataset(path="some/path", feature_names=["f1", "f2"])

# 2. EDGE TEST CASES

def test_empty_dataframe():
    # Test with empty dataframe
    df = pd.DataFrame()
    entity = _EntityEvaluationDataset(df)
    eval_ds = EvaluationDataset(entity)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 799μs -> 253μs (215% faster)

def test_none_name_and_digest():
    # Test with None for name and digest
    df = pd.DataFrame({"a": [1]})
    entity = _EntityEvaluationDataset(df, name=None, digest=None)
    eval_ds = EvaluationDataset(entity)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 1.29ms -> 481μs (167% faster)

def test_dataframe_with_special_types():
    # Test with dataframe containing various types
    df = pd.DataFrame({
        "int": [1, 2],
        "float": [1.1, 2.2],
        "str": ["foo", "bar"],
        "bool": [True, False],
        "none": [None, None]
    })
    entity = _EntityEvaluationDataset(df, name="special", digest="types")
    eval_ds = EvaluationDataset(entity)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 1.40ms -> 1.38ms (1.57% faster)

def test_databricks_dataset_with_empty_dataframe():
    # Databricks dataset with empty dataframe
    df = pd.DataFrame()
    dbr_ds = DatabricksDataset(df)
    eval_ds = EvaluationDataset(dbr_ds)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 769μs -> 253μs (204% faster)


def test_path_is_empty_string():
    # Path is empty string
    df = pd.DataFrame({"a": [1]})
    entity = _EntityEvaluationDataset(df)
    eval_ds = EvaluationDataset(entity)
    pyfunc_ds = eval_ds.to_evaluation_dataset(path="")

def test_large_string_in_name_and_digest():
    # Large string values for name and digest
    large_str = "x" * 500
    df = pd.DataFrame({"a": [1, 2]})
    entity = _EntityEvaluationDataset(df, name=large_str, digest=large_str)
    eval_ds = EvaluationDataset(entity)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 491μs -> 487μs (0.683% faster)

# 3. LARGE SCALE TEST CASES

def test_large_dataframe_conversion():
    # Test with a large dataframe (1000 rows, 10 columns)
    data = {f"col{i}": list(range(1000)) for i in range(10)}
    df = pd.DataFrame(data)
    entity = _EntityEvaluationDataset(df, name="large_ds", digest="digest_large")
    eval_ds = EvaluationDataset(entity)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 2.51ms -> 2.43ms (3.28% faster)

def test_large_databricks_dataset_conversion():
    # Large Databricks dataset
    data = {f"c{i}": [str(j) for j in range(1000)] for i in range(5)}
    df = pd.DataFrame(data)
    dbr_ds = DatabricksDataset(df, name="big_dbr", digest="digest_dbr")
    eval_ds = EvaluationDataset(dbr_ds)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds = codeflash_output # 6.32ms -> 6.14ms (3.07% faster)


def test_multiple_conversions_consistency():
    # Ensure repeated conversions are consistent
    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
    entity = _EntityEvaluationDataset(df, name="repeat", digest="repeat_digest")
    eval_ds = EvaluationDataset(entity)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds1 = codeflash_output # 590μs -> 577μs (2.19% faster)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds2 = codeflash_output # 426μs -> 414μs (2.93% faster)

def test_databricks_dataset_multiple_conversions_consistency():
    # Ensure repeated conversions are consistent for DatabricksDataset
    df = pd.DataFrame({"x": [10, 20, 30], "y": [40, 50, 60]})
    dbr_ds = DatabricksDataset(df, name="repeat_dbr", digest="digest_dbr")
    eval_ds = EvaluationDataset(dbr_ds)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds1 = codeflash_output # 576μs -> 571μs (0.911% faster)
    codeflash_output = eval_ds._to_pyfunc_dataset(); pyfunc_ds2 = codeflash_output # 428μs -> 414μs (3.20% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from mlflow.genai.datasets.evaluation_dataset import EvaluationDataset


# Minimal stub for pandas.DataFrame for testing without pandas
class DummyDataFrame:
    def __init__(self, data):
        self._data = data

    def __eq__(self, other):
        if not isinstance(other, DummyDataFrame):
            return False
        return self._data == other._data

    def __repr__(self):
        return f"DummyDataFrame({self._data})"

# Minimal stub for mlflow.entities.evaluation_dataset.EvaluationDataset
class EntityEvaluationDataset:
    def __init__(self, data, name=None, digest=None):
        self._data = data
        self.name = name
        self.digest = digest
        self._records = None

    def to_df(self):
        return DummyDataFrame(self._data)

    def __eq__(self, other):
        if not isinstance(other, EntityEvaluationDataset):
            return False
        return self._data == other._data and self.name == other.name and self.digest == other.digest

# Minimal stub for legacy EvaluationDataset (mlflow.data.evaluation_dataset.EvaluationDataset)
class LegacyEvaluationDataset:
    def __init__(self, data, path=None, feature_names=None, name=None, digest=None):
        self.data = data  # DummyDataFrame
        self.path = path
        self.feature_names = feature_names
        self.name = name
        self.digest = digest

    def __eq__(self, other):
        if not isinstance(other, LegacyEvaluationDataset):
            return False
        return (
            self.data == other.data and
            self.path == other.path and
            self.feature_names == other.feature_names and
            self.name == other.name and
            self.digest == other.digest
        )

    def __repr__(self):
        return (f"LegacyEvaluationDataset(data={self.data}, path={self.path}, "
                f"feature_names={self.feature_names}, name={self.name}, digest={self.digest})")

# Minimal stub for databricks.agents.datasets.Dataset
class DatabricksDataset:
    def __init__(self, data, name=None, digest=None):
        self._data = data
        self.name = name
        self.digest = digest

    def to_df(self):
        return DummyDataFrame(self._data)

    def __eq__(self, other):
        if not isinstance(other, DatabricksDataset):
            return False
        return self._data == other._data and self.name == other.name and self.digest == other.digest

# ------------------ UNIT TESTS ------------------

# Basic Test Cases






def test_pyfunc_dataset_with_missing_attributes():
    """Test __getattr__ raises AttributeError for missing attribute."""
    entity = EntityEvaluationDataset(data=[{"a": 1}])
    eval_ds = EvaluationDataset(entity)
    with pytest.raises(AttributeError):
        _ = eval_ds.non_existent_attribute

def test_pyfunc_dataset_private_attribute_access():
    """Test __getattr__ raises AttributeError for private attribute."""
    entity = EntityEvaluationDataset(data=[{"a": 1}])
    eval_ds = EvaluationDataset(entity)
    with pytest.raises(AttributeError):
        _ = eval_ds._private

def test_pyfunc_dataset_setattr_records():
    """Test __setattr__ propagates _records to entity dataset."""
    entity = EntityEvaluationDataset(data=[{"a": 1}])
    eval_ds = EvaluationDataset(entity)
    eval_ds._records = "my_records"

def test_pyfunc_dataset_eq_entity_and_eval():
    """Test __eq__ for EvaluationDataset and EntityEvaluationDataset."""
    entity = EntityEvaluationDataset(data=[{"a": 1}], name="foo", digest="abc")
    eval_ds = EvaluationDataset(entity)

def test_pyfunc_dataset_eq_databricks():
    """Test __eq__ for two EvaluationDataset objects wrapping DatabricksDataset."""
    db1 = DatabricksDataset(data=[{"x": 1}], name="foo", digest="abc")
    db2 = DatabricksDataset(data=[{"x": 1}], name="foo", digest="abc")
    eval_ds1 = EvaluationDataset(db1)
    eval_ds2 = EvaluationDataset(db2)

def test_pyfunc_dataset_eq_mixed_types():
    """Test __eq__ returns False for mixed types."""
    entity = EntityEvaluationDataset(data=[{"a": 1}], name="foo", digest="abc")
    db = DatabricksDataset(data=[{"a": 1}], name="foo", digest="abc")
    eval_ds_entity = EvaluationDataset(entity)
    eval_ds_db = EvaluationDataset(db)

# Large Scale Test Cases

from mlflow.genai.datasets.evaluation_dataset import EvaluationDataset
import pytest

def test_EvaluationDataset__to_pyfunc_dataset():
    with pytest.raises(AttributeError, match="'SymbolicInt'\\ object\\ has\\ no\\ attribute\\ 'to_df'"):
        EvaluationDataset._to_pyfunc_dataset(EvaluationDataset(0))

To edit these changes git checkout codeflash/optimize-EvaluationDataset._to_pyfunc_dataset-mhx2xczx and push.