⚡️ Speed up function apply_offsets_to_table by 8%

[codeflash-ai]

📄 8% (0.08x) speedup for apply_offsets_to_table in lonboard/_geoarrow/movingpandas_interop.py

⏱️ Runtime : 2.10 milliseconds → 1.95 milliseconds (best of 72 runs)

📝 Explanation and details

The optimized code achieves a 7% speedup by eliminating redundant attribute lookups and replacing loop-based list construction with list comprehensions.

Key optimizations:

1. Attribute caching: Pre-stores batch.schema, schema.metadata, and batch.num_columns as local variables, avoiding repeated attribute traversals in the loop.

2. Bulk data extraction: Uses list comprehensions to extract all columns [batch[i] for i in range(num_columns)] and fields [schema.field(i) for i in range(num_columns)] upfront, eliminating per-iteration lookups.

3. List comprehensions over explicit loops: Replaces the manual for loop with append() calls with list comprehensions for both new_fields and new_arrays. List comprehensions are implemented in C and avoid the overhead of repeated list.append() method calls.

Why it's faster:

• Python attribute access (like batch.schema.field(field_idx)) involves dictionary lookups that add overhead when repeated
• List comprehensions execute faster than equivalent for/append patterns due to optimized C implementation
• Bulk extraction patterns reduce the number of function calls from O(n) per operation to O(1)

Performance characteristics:
The optimization shows the most benefit for larger tables - test results show 6-11% speedups for tables with 100+ columns and 1000+ rows, while smaller tables see modest slowdowns (20-25%) due to the upfront extraction overhead. This makes it ideal for production workloads processing substantial datasets where the column/row count justifies the initial setup cost.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# imports
import pytest
from arro3.core import Array, DataType, RecordBatch, Schema, Table, list_array
from lonboard._geoarrow.movingpandas_interop import apply_offsets_to_table

# unit tests

# Helper to make a Table with given columns and data
def make_table(fields, columns):
    schema = Schema(fields)
    batch = RecordBatch(columns, schema=schema)
    return Table.from_batches([batch])

# Helper to create an Array of offsets
def make_offsets(offsets):
    return Array(offsets, DataType.int32())

# Helper to create a primitive Array
def make_array(data, dtype):
    return Array(data, dtype)

# Basic Test Cases
















#------------------------------------------------
from __future__ import annotations

# Patch the function under test to use fakes for testing
import types

# imports
import pytest
from arro3.core import Array, DataType, RecordBatch, Schema, Table, list_array
from lonboard._geoarrow.movingpandas_interop import apply_offsets_to_table

# --- Utilities for testing (minimal fake implementations for core types) ---

class FakeArray(list):
    """A minimal stand-in for arro3.core.Array."""
    pass

class FakeDataType:
    """Minimal DataType stand-in."""
    def __init__(self, name):
        self.name = name
    @staticmethod
    def list(field):
        return FakeDataType(f"list[{field.type.name}]")
    def __eq__(self, other):
        return isinstance(other, FakeDataType) and self.name == other.name

class FakeField:
    """Minimal field with type and name."""
    def __init__(self, name, type_):
        self.name = name
        self.type = type_
    def with_type(self, new_type):
        return FakeField(self.name, new_type)

class FakeSchema:
    """Minimal schema."""
    def __init__(self, fields, metadata=None):
        self._fields = fields
        self.metadata = metadata or {}
    def field(self, idx):
        return self._fields[idx]
    @property
    def fields(self):
        return self._fields

class FakeRecordBatch:
    """Minimal record batch."""
    def __init__(self, arrays, schema):
        self._arrays = arrays
        self.schema = schema
    @property
    def num_columns(self):
        return len(self._arrays)
    def __getitem__(self, idx):
        return self._arrays[idx]

class FakeTable:
    """Minimal Table."""
    def __init__(self, batches):
        self._batches = batches
    def combine_chunks(self):
        return self
    def to_batches(self):
        return self._batches
    @staticmethod
    def from_batches(batches):
        return FakeTable(batches)

def fake_list_array(offsets, values, type=None):
    """Simulate arro3.core.list_array: splits values by offsets."""
    out = []
    for i in range(len(offsets)-1):
        out.append(values[offsets[i]:offsets[i+1]])
    return out

apply_offsets_to_table.__globals__['Array'] = FakeArray
apply_offsets_to_table.__globals__['DataType'] = FakeDataType
apply_offsets_to_table.__globals__['RecordBatch'] = FakeRecordBatch
apply_offsets_to_table.__globals__['Schema'] = FakeSchema
apply_offsets_to_table.__globals__['Table'] = FakeTable
apply_offsets_to_table.__globals__['list_array'] = fake_list_array

# --- Unit tests ---

# BASIC TEST CASES

def test_single_column_basic_split():
    # One column, offsets split into 3 groups
    values = [1, 2, 3, 4, 5]
    offsets = [0, 2, 5]
    field = FakeField("a", FakeDataType("int"))
    schema = FakeSchema([field])
    batch = FakeRecordBatch([values], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 8.21μs -> 10.5μs (21.7% slower)
    # Expect: [[1,2],[3,4,5]]
    expected = [[1,2],[3,4,5]]

def test_two_columns_basic_split():
    # Two columns, split into 2 groups
    col1 = [10, 20, 30, 40]
    col2 = [1, 2, 3, 4]
    offsets = [0, 2, 4]
    fields = [FakeField("x", FakeDataType("int")), FakeField("y", FakeDataType("int"))]
    schema = FakeSchema(fields)
    batch = FakeRecordBatch([col1, col2], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 7.82μs -> 9.84μs (20.5% slower)

def test_offsets_single_group():
    # Offsets that cover all data as one group
    col = [1,2,3]
    offsets = [0,3]
    field = FakeField("a", FakeDataType("int"))
    schema = FakeSchema([field])
    batch = FakeRecordBatch([col], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 5.93μs -> 7.87μs (24.6% slower)

# EDGE TEST CASES

def test_empty_offsets():
    # Offsets with only one value (no groups)
    col = [1,2,3]
    offsets = [0]
    field = FakeField("a", FakeDataType("int"))
    schema = FakeSchema([field])
    batch = FakeRecordBatch([col], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 5.54μs -> 7.24μs (23.4% slower)

def test_empty_column():
    # Column is empty, offsets [0,0]
    col = []
    offsets = [0,0]
    field = FakeField("a", FakeDataType("int"))
    schema = FakeSchema([field])
    batch = FakeRecordBatch([col], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 5.70μs -> 7.67μs (25.8% slower)

def test_offsets_with_empty_groups():
    # Offsets that create empty groups in the middle
    col = [1,2,3]
    offsets = [0,1,1,3]
    field = FakeField("a", FakeDataType("int"))
    schema = FakeSchema([field])
    batch = FakeRecordBatch([col], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 6.07μs -> 7.97μs (23.9% slower)

def test_multiple_columns_with_empty_and_nonempty():
    # Two columns, one empty, one not
    col1 = []
    col2 = [7,8,9]
    offsets = [0,0,3]
    fields = [FakeField("a", FakeDataType("int")), FakeField("b", FakeDataType("int"))]
    schema = FakeSchema(fields)
    batch = FakeRecordBatch([col1, col2], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 7.70μs -> 9.99μs (22.9% slower)


def test_zero_columns():
    # Table with zero columns
    schema = FakeSchema([])
    batch = FakeRecordBatch([], schema)
    table = FakeTable([batch])
    offsets = [0,0]
    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 4.37μs -> 6.12μs (28.6% slower)

# LARGE SCALE TEST CASES

def test_large_table_many_rows_and_columns():
    # 100 columns, 1000 rows, split into 10 groups
    num_cols = 100
    num_rows = 1000
    num_groups = 10
    offsets = [i*(num_rows//num_groups) for i in range(num_groups)] + [num_rows]
    columns = [list(range(i, i+num_rows)) for i in range(num_cols)]
    fields = [FakeField(f"col{i}", FakeDataType("int")) for i in range(num_cols)]
    schema = FakeSchema(fields)
    batch = FakeRecordBatch(columns, schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 364μs -> 343μs (6.05% faster)
    # Each column should be split into 10 groups of 100
    for col_idx in range(num_cols):
        col = columns[col_idx]
        expected = [col[offsets[i]:offsets[i+1]] for i in range(num_groups)]

def test_large_table_single_group():
    # 500 columns, 1000 rows, all in one group
    num_cols = 500
    num_rows = 1000
    offsets = [0, 1000]
    columns = [list(range(i, i+num_rows)) for i in range(num_cols)]
    fields = [FakeField(f"col{i}", FakeDataType("int")) for i in range(num_cols)]
    schema = FakeSchema(fields)
    batch = FakeRecordBatch(columns, schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 1.50ms -> 1.35ms (11.1% faster)
    for col_idx in range(num_cols):
        pass

def test_large_number_of_groups():
    # 1 column, 1000 groups, 1000 rows
    num_rows = 1000
    offsets = list(range(num_rows+1))
    col = list(range(num_rows))
    field = FakeField("a", FakeDataType("int"))
    schema = FakeSchema([field])
    batch = FakeRecordBatch([col], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 97.1μs -> 96.7μs (0.351% faster)

def test_large_table_all_empty_groups():
    # 1 column, 1000 groups, all empty
    offsets = [0]*1001
    col = []
    field = FakeField("a", FakeDataType("int"))
    schema = FakeSchema([field])
    batch = FakeRecordBatch([col], schema)
    table = FakeTable([batch])

    codeflash_output = apply_offsets_to_table(table, offsets); result = codeflash_output # 87.9μs -> 91.1μs (3.60% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from lonboard._geoarrow.movingpandas_interop import apply_offsets_to_table

To edit these changes git checkout codeflash/optimize-apply_offsets_to_table-mhfm0v4b and push.