⚡️ Speed up function _apply_transforms by 42%

[codeflash-ai]

📄 42% (0.42x) speedup for _apply_transforms in marimo/_plugins/ui/_impl/dataframes/transforms/apply.py

⏱️ Runtime : 15.2 microseconds → 10.8 microseconds (best of 48 runs)

📝 Explanation and details

The optimized code achieves a 41% speedup by replacing a linear chain of 12 if statements with a single dictionary lookup, eliminating expensive repeated identity comparisons.

Key Optimizations:

1. Dictionary Dispatch: The original code used a chain of if transform.type is TransformType.X statements that required up to 12 identity comparisons per call. The optimized version uses a precomputed dictionary _transform_type_to_handler_method that provides O(1) lookup time regardless of transform type.

2. Reduced Branching: Instead of 12 conditional branches, there's now just one dictionary lookup followed by a single getattr() call. This eliminates the CPU pipeline stalls caused by unpredictable branching.

3. Attribute Caching: The transforms.transforms list is cached as transforms_list to avoid repeated attribute lookups in the loop.

Performance Impact:

• The line profiler shows the _handle function's total time dropped from 211µs to 107µs (49% faster)
• The dictionary lookup (method_name = _transform_type_to_handler_method.get(transform.type)) takes only 25µs vs the original chain of comparisons taking 140µs
• Test cases with unknown transform types see dramatic speedups (57-63% faster) due to faster failure detection

Hot Path Benefits:
Based on the function references, _apply_transforms is called from the apply() method in dataframe transformation pipelines, potentially processing multiple transforms per operation. This optimization will have compounding benefits when processing batches of transforms, as each _handle call is now significantly faster.

The optimization is particularly effective for transforms later in the enum sequence (like UNIQUE, EXPAND_DICT) that previously required checking all preceding conditions.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

🌀 Generated Regression Tests and Runtime

from typing import Any, List, Optional

# imports
import pytest
from marimo._plugins.ui._impl.dataframes.transforms.apply import \
    _apply_transforms

# --- Minimal stubs for types and handlers ---


# Enum for TransformType
class TransformType:
    COLUMN_CONVERSION = "COLUMN_CONVERSION"
    RENAME_COLUMN = "RENAME_COLUMN"
    SORT_COLUMN = "SORT_COLUMN"
    FILTER_ROWS = "FILTER_ROWS"
    GROUP_BY = "GROUP_BY"
    AGGREGATE = "AGGREGATE"
    SELECT_COLUMNS = "SELECT_COLUMNS"
    SHUFFLE_ROWS = "SHUFFLE_ROWS"
    SAMPLE_ROWS = "SAMPLE_ROWS"
    EXPLODE_COLUMNS = "EXPLODE_COLUMNS"
    EXPAND_DICT = "EXPAND_DICT"
    UNIQUE = "UNIQUE"

# Transform object
class Transform:
    def __init__(self, type_: str, payload: Optional[Any] = None):
        self.type = type_
        self.payload = payload

# Transformations object
class Transformations:
    def __init__(self, transforms: Optional[List[Transform]] = None):
        self.transforms = transforms or []

# Minimal TransformHandler with logging
class DummyHandler:
    def __init__(self):
        self.calls = []

    @staticmethod
    def handle_column_conversion(df, transform):
        # Simulate transformation
        df["column_conversion"] = True
        return df

    @staticmethod
    def handle_rename_column(df, transform):
        df["rename_column"] = True
        return df

    @staticmethod
    def handle_sort_column(df, transform):
        df["sort_column"] = True
        return df

    @staticmethod
    def handle_filter_rows(df, transform):
        df["filter_rows"] = True
        return df

    @staticmethod
    def handle_group_by(df, transform):
        df["group_by"] = True
        return df

    @staticmethod
    def handle_aggregate(df, transform):
        df["aggregate"] = True
        return df

    @staticmethod
    def handle_select_columns(df, transform):
        df["select_columns"] = True
        return df

    @staticmethod
    def handle_shuffle_rows(df, transform):
        df["shuffle_rows"] = True
        return df

    @staticmethod
    def handle_sample_rows(df, transform):
        df["sample_rows"] = True
        return df

    @staticmethod
    def handle_explode_columns(df, transform):
        df["explode_columns"] = True
        return df

    @staticmethod
    def handle_expand_dict(df, transform):
        df["expand_dict"] = True
        return df

    @staticmethod
    def handle_unique(df, transform):
        df["unique"] = True
        return df
from marimo._plugins.ui._impl.dataframes.transforms.apply import \
    _apply_transforms

# --- Unit tests ---

# BASIC TEST CASES

def test_apply_transforms_no_transforms_returns_original():
    # Scenario: No transforms should return the original df
    df = {"data": 123}
    handler = DummyHandler()
    transformations = Transformations([])
    codeflash_output = _apply_transforms(df.copy(), handler, transformations); result = codeflash_output # 514ns -> 513ns (0.195% faster)






def test_apply_transforms_with_unknown_transform_type_raises():
    # Scenario: Unknown transform type should raise AssertionError via assert_never
    df = {}
    handler = DummyHandler()
    transformations = Transformations([Transform("UNKNOWN_TYPE")])
    with pytest.raises(AssertionError):
        _apply_transforms(df.copy(), handler, transformations) # 4.62μs -> 2.94μs (57.2% faster)

from typing import Any, List

# imports
import pytest  # used for our unit tests
from marimo._plugins.ui._impl.dataframes.transforms.apply import \
    _apply_transforms

# --- Minimal stubs for types and handlers ---

# Simulate TransformType as an Enum
class TransformType:
    COLUMN_CONVERSION = "COLUMN_CONVERSION"
    RENAME_COLUMN = "RENAME_COLUMN"
    SORT_COLUMN = "SORT_COLUMN"
    FILTER_ROWS = "FILTER_ROWS"
    GROUP_BY = "GROUP_BY"
    AGGREGATE = "AGGREGATE"
    SELECT_COLUMNS = "SELECT_COLUMNS"
    SHUFFLE_ROWS = "SHUFFLE_ROWS"
    SAMPLE_ROWS = "SAMPLE_ROWS"
    EXPLODE_COLUMNS = "EXPLODE_COLUMNS"
    EXPAND_DICT = "EXPAND_DICT"
    UNIQUE = "UNIQUE"

# Simulate Transform object
class Transform:
    def __init__(self, type_: str, **kwargs):
        self.type = type_
        for k, v in kwargs.items():
            setattr(self, k, v)

# Simulate Transformations object
class Transformations:
    def __init__(self, transforms: List[Transform]):
        self.transforms = transforms

# Minimal TransformHandler implementation
class DummyHandler:
    # Each handler returns a tuple with the operation name and arguments for test verification
    @staticmethod
    def handle_column_conversion(df, transform):
        return df + [("column_conversion", transform.__dict__)]
    @staticmethod
    def handle_rename_column(df, transform):
        return df + [("rename_column", transform.__dict__)]
    @staticmethod
    def handle_sort_column(df, transform):
        return df + [("sort_column", transform.__dict__)]
    @staticmethod
    def handle_filter_rows(df, transform):
        return df + [("filter_rows", transform.__dict__)]
    @staticmethod
    def handle_group_by(df, transform):
        return df + [("group_by", transform.__dict__)]
    @staticmethod
    def handle_aggregate(df, transform):
        return df + [("aggregate", transform.__dict__)]
    @staticmethod
    def handle_select_columns(df, transform):
        return df + [("select_columns", transform.__dict__)]
    @staticmethod
    def handle_shuffle_rows(df, transform):
        return df + [("shuffle_rows", transform.__dict__)]
    @staticmethod
    def handle_sample_rows(df, transform):
        return df + [("sample_rows", transform.__dict__)]
    @staticmethod
    def handle_explode_columns(df, transform):
        return df + [("explode_columns", transform.__dict__)]
    @staticmethod
    def handle_expand_dict(df, transform):
        return df + [("expand_dict", transform.__dict__)]
    @staticmethod
    def handle_unique(df, transform):
        return df + [("unique", transform.__dict__)]
from marimo._plugins.ui._impl.dataframes.transforms.apply import \
    _apply_transforms

# --- Unit tests ---

# 1. Basic Test Cases

def test_no_transforms_returns_input():
    # No transforms: should return original df
    df = ["original"]
    handler = DummyHandler()
    transforms = Transformations([])
    codeflash_output = _apply_transforms(df, handler, transforms); result = codeflash_output # 539ns -> 542ns (0.554% slower)




def test_empty_df_and_no_transforms():
    # Edge: empty df and no transforms
    df = []
    handler = DummyHandler()
    transforms = Transformations([])
    codeflash_output = _apply_transforms(df, handler, transforms); result = codeflash_output # 496ns -> 502ns (1.20% slower)


def test_transform_with_unexpected_type_raises():
    # Edge: unknown transform type should raise AssertionError
    df = []
    handler = DummyHandler()
    t = Transform("UNKNOWN_TYPE")
    transforms = Transformations([t])
    with pytest.raises(AssertionError) as excinfo:
        _apply_transforms(df, handler, transforms) # 4.61μs -> 2.83μs (62.8% faster)

def test_transform_with_none_type_raises():
    # Edge: None type should raise AssertionError
    df = []
    handler = DummyHandler()
    t = Transform(None)
    transforms = Transformations([t])
    with pytest.raises(AssertionError) as excinfo:
        _apply_transforms(df, handler, transforms) # 4.03μs -> 2.89μs (39.2% faster)




def test_transformations_is_none():
    # Edge: transforms.transforms is None
    class BadTransformations:
        def __init__(self, transforms):
            self.transforms = transforms
    df = []
    handler = DummyHandler()
    transforms = BadTransformations(None)
    codeflash_output = _apply_transforms(df, handler, transforms); result = codeflash_output # 425ns -> 528ns (19.5% slower)

# 3. Large Scale Test Cases

To edit these changes git checkout codeflash/optimize-_apply_transforms-mhwuuqp1 and push.