⚡️ Speed up function get_default_table_page_size by 95%

[codeflash-ai]

📄 95% (0.95x) speedup for get_default_table_page_size in marimo/_plugins/ui/_impl/table.py

⏱️ Runtime : 427 microseconds → 219 microseconds (best of 250 runs)

📝 Explanation and details

The optimization achieves a 94% speedup by eliminating expensive exception handling in the common case through direct context access instead of function indirection.

Key Changes:

1. Direct context access: Instead of calling get_context() which checks _THREAD_LOCAL_CONTEXT.runtime_context and raises an exception, the optimized code directly accesses _THREAD_LOCAL_CONTEXT.runtime_context
2. Exception elimination: Replaced try/except block with a simple if ctx is None check
3. Added import: Imported _THREAD_LOCAL_CONTEXT directly to enable the optimization

Why This is Faster:

• Exception handling overhead: Python's try/except mechanism is expensive even when no exception occurs. The line profiler shows the original get_context() call taking 3.34ms (80.5% of total time) vs only 311μs (39.4%) for direct access
• Function call elimination: Removes the overhead of calling get_context() entirely in the optimized version (only 6 calls vs 1175 in original)
• Simpler control flow: Direct attribute access and null check is much faster than exception-based control flow

Performance Benefits:
The test results show consistent 95-139% speedup across all scenarios:

• Context not initialized: 126% faster (most common case)
• Various configured page sizes: 102-139% faster
• Large-scale operations: 90-99% faster for bulk operations

This optimization is particularly effective because it targets the hot path where context checking happens frequently, converting expensive exception handling into fast conditional logic while maintaining identical behavior.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
_plugins/ui/_impl/test_table.py::test_default_table_page_size	1.55μs	674ns	130%✅

🌀 Generated Regression Tests and Runtime

import pytest
from marimo._plugins.ui._impl.table import get_default_table_page_size

function to test

class ContextNotInitializedError(Exception):
pass

class DummyContext:
def init(self, page_size):
self.marimo_config = {
"display": {
"default_table_page_size": page_size
}
}

Simulate thread-local context for testing

class _ThreadLocalContext:
def init(self):
self.runtime_context = None

_THREAD_LOCAL_CONTEXT = _ThreadLocalContext()
from marimo._plugins.ui._impl.table import get_default_table_page_size

unit tests

--- Basic Test Cases ---

def test_returns_default_when_context_not_initialized():
"""Should return 10 if context is not initialized."""
_THREAD_LOCAL_CONTEXT.runtime_context = None
codeflash_output = get_default_table_page_size() # 1.83μs -> 809ns (126% faster)

def test_returns_configured_page_size():
"""Should return the value from context config if present."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=25)
codeflash_output = get_default_table_page_size() # 1.60μs -> 671ns (139% faster)

def test_returns_configured_page_size_with_other_values():
"""Should return the value from context config for different values."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=50)
codeflash_output = get_default_table_page_size() # 1.67μs -> 730ns (128% faster)

--- Edge Test Cases ---

def test_configured_page_size_zero():
"""Should handle zero page size."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=0)
codeflash_output = get_default_table_page_size() # 1.49μs -> 738ns (102% faster)

def test_configured_page_size_negative():
"""Should handle negative page size."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=-5)
codeflash_output = get_default_table_page_size() # 1.46μs -> 682ns (115% faster)

def test_configured_page_size_one():
"""Should handle minimum positive page size."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=1)
codeflash_output = get_default_table_page_size() # 1.44μs -> 677ns (113% faster)

def test_configured_page_size_large_int():
"""Should handle very large page size."""
large_value = 10**6
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=large_value)
codeflash_output = get_default_table_page_size() # 1.34μs -> 648ns (106% faster)

def test_config_page_size_is_none():
"""Should return None if config value is None."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=None)
codeflash_output = get_default_table_page_size() # 2.31μs -> 992ns (133% faster)

def test_config_page_size_is_float():
"""Should return float if config value is float."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=15.5)
codeflash_output = get_default_table_page_size() # 1.63μs -> 740ns (120% faster)

def test_config_page_size_is_string():
"""Should return string if config value is string."""
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size="20")
codeflash_output = get_default_table_page_size() # 1.53μs -> 694ns (120% faster)

--- Large Scale Test Cases ---

def test_many_context_switches():
"""Should always return correct value after many context changes."""
for i in range(100):
_THREAD_LOCAL_CONTEXT.runtime_context = DummyContext(page_size=i)
codeflash_output = get_default_table_page_size() # 35.4μs -> 18.6μs (90.5% faster)
# Reset to None for next test
_THREAD_LOCAL_CONTEXT.runtime_context = None
codeflash_output = get_default_table_page_size() # 322ns -> 174ns (85.1% faster)

def test_large_context_config_dict():
"""Should work with large config dicts."""
big_dict = {"display": {"default_table_page_size": 777}}
# Add lots of extra keys
for i in range(900):
big_dict[f"extra_{i}"] = i
class BigContext:
def init(self):
self.marimo_config = big_dict
_THREAD_LOCAL_CONTEXT.runtime_context = BigContext()
codeflash_output = get_default_table_page_size() # 1.58μs -> 721ns (119% faster)

def test_large_number_of_contexts():
"""Should work with many different context objects."""
contexts = [DummyContext(page_size=i) for i in range(1000)]
for i, ctx in enumerate(contexts):
_THREAD_LOCAL_CONTEXT.runtime_context = ctx
codeflash_output = get_default_table_page_size() # 327μs -> 170μs (92.2% 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 marimo._plugins.ui._impl.table import get_default_table_page_size

---- Function to test ----

Simulate ContextNotInitializedError and get_context for testing purposes

class ContextNotInitializedError(Exception):
"""Custom exception to simulate context not being initialized."""
pass

Simulated global context object for thread-local storage

class DummyContext:
def init(self, config):
self.marimo_config = config

We'll use a global variable to simulate thread-local context for tests

_GLOBAL_CONTEXT = {"initialized": False, "context": None}
from marimo._plugins.ui._impl.table import get_default_table_page_size

---- Unit tests ----

Helper to set up the global context for tests

def setup_context(page_size):
_GLOBAL_CONTEXT["initialized"] = True
_GLOBAL_CONTEXT["context"] = DummyContext(
{"display": {"default_table_page_size": page_size}}
)

def teardown_context():
_GLOBAL_CONTEXT["initialized"] = False
_GLOBAL_CONTEXT["context"] = None

------------------- Basic Test Cases -------------------

def test_returns_default_when_context_not_initialized():
"""Should return 10 when context is not initialized."""
teardown_context()
codeflash_output = get_default_table_page_size() # 1.80μs -> 821ns (119% faster)

def test_returns_default_when_context_is_none():
"""Should return 10 when context is explicitly set to None."""
_GLOBAL_CONTEXT["initialized"] = True
_GLOBAL_CONTEXT["context"] = None
codeflash_output = get_default_table_page_size() # 1.34μs -> 620ns (116% faster)
teardown_context()

def test_returns_configured_page_size():
"""Should return the configured page size from context."""
setup_context(25)
codeflash_output = get_default_table_page_size() # 1.44μs -> 609ns (136% faster)
teardown_context()

def test_returns_configured_page_size_for_minimum_valid():
"""Should return the minimum valid page size (e.g., 1)."""
setup_context(1)
codeflash_output = get_default_table_page_size() # 1.33μs -> 673ns (97.0% faster)
teardown_context()

def test_returns_configured_page_size_for_typical_value():
"""Should return a typical page size (e.g., 50)."""
setup_context(50)
codeflash_output = get_default_table_page_size() # 1.41μs -> 612ns (130% faster)
teardown_context()

------------------- Edge Test Cases -------------------

def test_returns_configured_page_size_for_zero():
"""Edge: Should handle zero as a page size (if allowed)."""
setup_context(0)
codeflash_output = get_default_table_page_size() # 1.38μs -> 669ns (107% faster)
teardown_context()

def test_returns_configured_page_size_for_large_value():
"""Edge: Should handle a very large page size."""
large_value = 999
setup_context(large_value)
codeflash_output = get_default_table_page_size() # 1.47μs -> 686ns (114% faster)
teardown_context()

def test_returns_configured_page_size_for_negative_value():
"""Edge: Should handle negative page size (if allowed)."""
setup_context(-5)
codeflash_output = get_default_table_page_size() # 1.46μs -> 685ns (113% faster)
teardown_context()

def test_returns_configured_page_size_for_nonstandard_value():
"""Edge: Should handle nonstandard but valid page size."""
setup_context(17)
codeflash_output = get_default_table_page_size() # 1.50μs -> 704ns (112% faster)
teardown_context()

------------------- Large Scale Test Cases -------------------

@pytest.mark.parametrize("page_size", list(range(1, 1001, 100)))
def test_large_scale_page_sizes(page_size):
"""Test function with a range of large page sizes."""
setup_context(page_size)
codeflash_output = get_default_table_page_size() # 14.2μs -> 6.67μs (112% faster)
teardown_context()

def test_large_scale_randomized_page_sizes():
"""Test function with many random page sizes under 1000."""
import random
for _ in range(20): # 20 random values for scalability
page_size = random.randint(1, 1000)
setup_context(page_size)
codeflash_output = get_default_table_page_size() # 8.09μs -> 4.14μs (95.5% faster)
teardown_context()

def test_large_scale_all_defaults():
"""Test function returns 10 for many uninitialized contexts."""
teardown_context()
for _ in range(20):
codeflash_output = get_default_table_page_size() # 7.63μs -> 3.84μs (99.0% faster)

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
from marimo._plugins.ui._impl.table import get_default_table_page_size

def test_get_default_table_page_size():
get_default_table_page_size()

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_bps3n5s8/tmpif_p40vd/test_concolic_coverage.py::test_get_default_table_page_size	1.34μs	625ns	115%✅

To edit these changes git checkout codeflash/optimize-get_default_table_page_size-mhvckegp and push.