Chain-of-Table: Evolving Tables in the Reasoning Chain for Table Understanding { }

| country | cyclist_count |
|---------|---------------|
| USA     | 2             |
| Canada  | 1             |
| Spain   | 1             |

import pandas as pd
from typing import List, Tuple, Dict, Any
import re

class TableOperation:
    """Base class for table operations."""
    def apply(self, table: pd.DataFrame) -> pd.DataFrame:
        raise NotImplementedError

class AddColumnOp(TableOperation):
    """Add a new column by extracting from an existing column."""
    def __init__(self, new_col_name: str, source_col: str, extraction_rule: str):
        self.new_col_name = new_col_name
        self.source_col = source_col
        self.extraction_rule = extraction_rule
    
    def apply(self, table: pd.DataFrame) -> pd.DataFrame:
        # Example: extraction_rule = "split_right_comma" 
        # means split by comma and take the right side
        result = table.copy()
        if "split_right_comma" in self.extraction_rule:
            result[self.new_col_name] = result[self.source_col].str.split(',').str[-1].str.strip()
        elif "split_left_comma" in self.extraction_rule:
            result[self.new_col_name] = result[self.source_col].str.split(',').str[0].str.strip()
        return result
    
    def __repr__(self):
        return f"ADD_COLUMN({self.new_col_name}, {self.source_col}, {self.extraction_rule})"

class SelectRowsOp(TableOperation):
    """Select specific rows by index."""
    def __init__(self, indices: List[int]):
        self.indices = indices
    
    def apply(self, table: pd.DataFrame) -> pd.DataFrame:
        return table.iloc[self.indices].reset_index(drop=True)
    
    def __repr__(self):
        return f"SELECT_ROWS({self.indices})"

class FilterOp(TableOperation):
    """Filter rows based on a condition."""
    def __init__(self, column: str, condition: str, value: Any):
        self.column = column
        self.condition = condition  # e.g., "<=", "==", ">", "in"
        self.value = value
    
    def apply(self, table: pd.DataFrame) -> pd.DataFrame:
        result = table.copy()
        if self.condition == "<=":
            result = result[result[self.column] <= self.value]
        elif self.condition == "==":
            result = result[result[self.column] == self.value]
        elif self.condition == ">":
            result = result[result[self.column] > self.value]
        elif self.condition == "in":
            result = result[result[self.column].isin(self.value)]
        return result.reset_index(drop=True)
    
    def __repr__(self):
        return f"FILTER({self.column}, {self.condition}, {self.value})"

class AggregateOp(TableOperation):
    """Aggregate rows by grouping and applying a reduction function."""
    def __init__(self, group_by_col: str, agg_func: str, output_col: str):
        self.group_by_col = group_by_col
        self.agg_func = agg_func  # e.g., "count", "sum", "mean"
        self.output_col = output_col
    
    def apply(self, table: pd.DataFrame) -> pd.DataFrame:
        result = table.copy()
        if self.agg_func == "count":
            grouped = result.groupby(self.group_by_col).size().reset_index(name=self.output_col)
        elif self.agg_func == "sum":
            # Assume we're summing all numeric columns except the grouping column
            numeric_cols = result.select_dtypes(include=['number']).columns.tolist()
            if self.group_by_col in numeric_cols:
                numeric_cols.remove(self.group_by_col)
            grouped = result.groupby(self.group_by_col)[numeric_cols].sum().reset_index()
            grouped = grouped.rename(columns={numeric_cols[0]: self.output_col} if numeric_cols else {})
        else:
            raise ValueError(f"Unknown aggregation function: {self.agg_func}")
        return grouped
    
    def __repr__(self):
        return f"AGGREGATE({self.group_by_col}, {self.agg_func}, {self.output_col})"

class SortOp(TableOperation):
    """Sort rows by a column."""
    def __init__(self, column: str, order: str = "ASC"):
        self.column = column
        self.order = order  # "ASC" or "DESC"
    
    def apply(self, table: pd.DataFrame) -> pd.DataFrame:
        ascending = (self.order == "ASC")
        return table.sort_values(by=self.column, ascending=ascending).reset_index(drop=True)
    
    def __repr__(self):
        return f"SORT({self.column}, {self.order})"

class ChainOfTableFramework:
    """Executes the Chain-of-Table reasoning pipeline."""
    
    def __init__(self, max_iterations: int = 10, min_table_size: int = 1):
        self.max_iterations = max_iterations
        self.min_table_size = min_table_size
        self.operation_chain: List[TableOperation] = []
    
    def add_operation(self, op: TableOperation) -> None:
        """Add an operation to the chain."""
        self.operation_chain.append(op)
    
    def get_operation_chain_str(self) -> str:
        """Return the operation chain as a human-readable string."""
        return " → ".join([str(op) for op in self.operation_chain])
    
    def execute(self, table: pd.DataFrame) -> pd.DataFrame:
        """Execute all operations in the chain, in order."""
        result = table.copy()
        print("Executing operation chain:")
        for i, op in enumerate(self.operation_chain):
            result = op.apply(result)
            print(f"  Step {i+1}: {op}")
            print(f"    Table shape: {result.shape}")
            if result.shape[0] <= 5:
                print(f"    Table:\n{result}")
        return result
    
    def reset(self) -> None:
        """Clear the operation chain."""
        self.operation_chain = []

# Example usage
if __name__ == "__main__":
    # Create a sample table of cyclist rankings
    data = {
        "Rank": [1, 2, 3, 4, 5],
        "Name_Country": ["John Smith, USA", "Marie Dupont, Canada", "Carlos Gomez, Spain", 
                         "Alice Johnson, USA", "Chen Wei, China"],
        "Points": [100, 95, 90, 85, 80]
    }
    table = pd.DataFrame(data)
    
    print("=" * 70)
    print("CHAIN-OF-TABLE EXAMPLE: Which country had the most cyclists in top 3?")
    print("=" * 70)
    print("\nOriginal table:")
    print(table)
    
    # Initialize framework
    cot = ChainOfTableFramework(max_iterations=10)
    
    # Step 1: Add a column to extract country from composite cell
    print("\n" + "=" * 70)
    print("STAGE 1: Operation Selection")
    print("=" * 70)
    print("\nLLM selects: ADD_COLUMN (needed to separate name from country)")
    cot.add_operation(AddColumnOp("Country", "Name_Country", "split_right_comma"))
    
    # Step 2: Filter to top 3 ranks
    print("\nLLM selects: FILTER (to keep only top 3 ranked cyclists)")
    cot.add_operation(FilterOp("Rank", "<=", 3))
    
    # Step 3: Aggregate by country
    print("\nLLM selects: AGGREGATE (to count cyclists per country)")
    cot.add_operation(AggregateOp("Country", "count", "Cyclist_Count"))
    
    # Step 4: Sort to find the maximum
    print("\nLLM selects: SORT (to order by count descending)")
    cot.add_operation(SortOp("Cyclist_Count", "DESC"))
    
    # Step 5: Done
    print("\nLLM selects: DONE (table is ready for final answer extraction)")
    
    print("\n" + "=" * 70)
    print("STAGE 3: Deterministic Execution")
    print("=" * 70)
    final_table = cot.execute(table)
    
    print("\n" + "=" * 70)
    print("STAGE 3b: Final Answer Extraction")
    print("=" * 70)
    print("\nFinal transformed table:")
    print(final_table)
    
    print("\nLLM generates final answer:")
    top_country = final_table.iloc[0]
    answer = f"{top_country['Country']} had the most cyclists in the top 3, with {int(top_country['Cyclist_Count'])} cyclist(s)."
    print(f"  → {answer}")
    
    print("\n" + "=" * 70)
    print("OPERATION CHAIN TRANSCRIPT")
    print("=" * 70)
    print(cot.get_operation_chain_str())

#!/usr/bin/env python3
"""
Minimal Chain-of-Table implementation demonstrating the core concept.
Requires: pip install openai python-dotenv
Environment: Set OPENAI_API_KEY
"""

import os
import json
from openai import OpenAI

# Initialize client from environment variable
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Example table: cyclist rankings with mixed-cell data
initial_table = [
    {"rank": 1, "name_country": "Alice Smith, USA", "points": 2500},
    {"rank": 2, "name_country": "Bob Johnson, Canada", "points": 2300},
    {"rank": 3, "name_country": "Carlos Ruiz, Spain", "points": 2100},
]

question = "Which country had the most cyclists finish in the top 3?"

# Stage 1: Ask LLM to identify the first transformation needed
prompt_stage1 = f"""You are analyzing a table to answer a question through iterative transformations.

Current table:
{json.dumps(initial_table, indent=2)}

Question: {question}

What is the first transformation needed? Choose from:
1. ADD_COLUMN(name, source_column, extraction_rule)
2. SELECT_ROWS(indices)
3. AGGREGATE(group_by_column, aggregation_func)
4. FILTER(column, condition)

Respond with only the operation in the format: OPERATION(param1, param2, ...)
Do not include explanations."""

response = client.messages.create(
    model="gpt-4o-mini",
    max_tokens=100,
    messages=[{"role": "user", "content": prompt_stage1}]
)

operation = response.content[0].text.strip()
print(f"Stage 1 - Identified operation: {operation}")

# Stage 2: Execute the transformation (simple simulation)
if operation.startswith("ADD_COLUMN"):
    # Parse operation and apply transformation
    for row in initial_table:
        name, country = row["name_country"].split(", ")
        row["country"] = country
        row["name"] = name
    print(f"Stage 2 - Added 'country' column via extraction")

# Stage 3: Ask LLM if more transformations are needed
prompt_stage3 = f"""Current table after transformations:
{json.dumps(initial_table[:3], indent=2)}

Question: {question}

Is the table now ready to answer the question? 
If yes, answer the question directly.
If no, specify the next transformation needed.

Respond with either:
- ANSWER: [your answer]
- OPERATION: [next transformation]"""

response_final = client.messages.create(
    model="gpt-4o-mini",
    max_tokens=200,
    messages=[{"role": "user", "content": prompt_stage3}]
)

print(f"Stage 3 - Final response:\n{response_final.content[0].text.strip()}")

#!/usr/bin/env python3
"""
Complete Chain-of-Table implementation with the full reasoning loop.
Includes operation registry, state tracking, and iterative refinement.
Requires: pip install openai python-dotenv
"""

import os
import json
import re
from typing import Any, Dict, List, Tuple, Optional
from dataclasses import dataclass, field
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

@dataclass
class TableState:
    """Maintains mutable table state and operation history for chain reasoning."""
    
    data: List[Dict[str, Any]]
    operation_chain: List[str] = field(default_factory=list)
    
    def columns(self) -> List[str]:
        """Return column names from first row."""
        return list(self.data[0].keys()) if self.data else []
    
    def to_string(self, max_rows: int = 10) -> str:
        """Serialize table to readable string, truncating if necessary."""
        if not self.data:
            return "Empty table"
        
        shown = self.data[:max_rows]
        truncated = len(self.data) > max_rows
        
        return json.dumps(shown, indent=2) + (
            f"\n... ({len(self.data) - max_rows} more rows)" if truncated else ""
        )


class OperationRegistry:
    """Registry of available table operations with argument parsing."""
    
    @staticmethod
    def add_column(
        table_state: TableState,
        new_column_name: str,
        source_column: str,
        extraction_rule: str
    ) -> TableState:
        """
        Add a new column by extracting data from an existing column.
        extraction_rule: plain English description (e.g., "split by comma and take the right side")
        """
        new_data = []
        for row in table_state.data:
            new_row = row.copy()
            source_value = str(row.get(source_column, ""))
            
            # Simple extraction rules
            if "split by comma" in extraction_rule.lower() and "right" in extraction_rule.lower():
                parts = source_value.split(", ")
                new_row[new_column_name] = parts[-1] if parts else source_value
            elif "split by comma" in extraction_rule.lower() and "left" in extraction_rule.lower():
                parts = source_value.split(", ")
                new_row[new_column_name] = parts[0] if parts else source_value
            elif "split by space" in extraction_rule.lower() and "first" in extraction_rule.lower():
                parts = source_value.split()
                new_row[new_column_name] = parts[0] if parts else source_value
            else:
                # Fallback: use source value as-is
                new_row[new_column_name] = source_value
            
            new_data.append(new_row)
        
        new_state = TableState(data=new_data)
        new_state.operation_chain = table_state.operation_chain + [
            f"ADD_COLUMN(name={new_column_name}, source={source_column})"
        ]
        return new_state
    
    @staticmethod
    def select_rows(
        table_state: TableState,
        indices: List[int]
    ) -> TableState:
        """Retain only rows at specified indices (0-indexed)."""
        selected = [table_state.data[i] for i in indices if i < len(table_state.data)]
        new_state = TableState(data=selected)
        new_state.operation_chain = table_state.operation_chain + [
            f"SELECT_ROWS(indices={indices})"
        ]
        return new_state
    
    @staticmethod
    def filter_rows(
        table_state: TableState,
        column: str,
        condition: str
    ) -> TableState:
        """
        Filter rows where column satisfies a condition.
        condition: plain English (e.g., "rank <= 3", "country == USA")
        """
        filtered = []
        for row in table_state.data:
            column_value = row.get(column)
            
            # Parse simple numeric conditions
            if "<=" in condition:
                parts = condition.split("<=")
                threshold = float(parts[1].strip())
                try:
                    if float(column_value) <= threshold:
                        filtered.append(row)
                except (ValueError, TypeError):
                    pass
            elif ">=" in condition:
                parts = condition.split(">=")
                threshold = float(parts[1].strip())
                try:
                    if float(column_value) >= threshold:
                        filtered.append(row)
                except (ValueError, TypeError):
                    pass
            elif "==" in condition:
                parts = condition.split("==")
                target = parts[1].strip().strip('"\'')
                if str(column_value).strip() == target:
                    filtered.append(row)
            elif "!=" in condition:
                parts = condition.split("!=")
                target = parts[1].strip().strip('"\'')
                if str(column_value).strip() != target:
                    filtered.append(row)
            else:
                # Fallback: keep all
                filtered.append(row)
        
        new_state = TableState(data=filtered)
        new_state.operation_chain = table_state.operation_chain + [
            f"FILTER(column={column}, condition={condition})"
        ]
        return new_state
    
    @staticmethod
    def aggregate(
        table_state: TableState,
        group_by_column: str,
        agg_func: str
    ) -> TableState:
        """
        Group by a column and apply an aggregation function.
        agg_func: "count", "sum", "avg", "max", "min"
        """
        groups: Dict[str, List[Dict]] = {}
        
        for row in table_state.data:
            key = str(row.get(group_by_column, "Unknown"))
            if key not in groups:
                groups[key] = []
            groups[key].append(row)
        
        aggregated = []
        for group_key, group_rows in groups.items():
            agg_row = {group_by_column: group_key}
            
            if agg_func.lower() == "count":
                agg_row["count"] = len(group_rows)
            elif agg_func.lower() == "sum":
                # Sum all numeric columns except group_by
                for col in group_rows[0].keys():
                    if col != group_by_column:
                        try:
                            agg_row[col] = sum(float(r.get(col, 0)) for r in group_rows)
                        except (ValueError, TypeError):
                            pass
            elif agg_func.lower() == "avg":
                # Average numeric columns
                for col in group_rows[0].keys():
                    if col != group_by_column:
                        try:
                            values = [float(r.get(col, 0)) for r in group_rows]
                            agg_row[col] = sum(values) / len(values) if values else 0
                        except (ValueError, TypeError):
                            pass
            elif agg_func.lower() == "max":
                for col in group_rows[0].keys():
                    if col != group_by_column:
                        try:
                            values = [float(r.get(col, 0)) for r in group_rows]
                            agg_row[col] = max(values) if values else 0
                        except (ValueError, TypeError):
                            pass
            elif agg_func.lower() == "min":
                for col in group_rows[0].keys():
                    if col != group_by_column:
                        try:
                            values = [float(r.get(col, 0)) for r in group_rows]
                            agg_row[col] = min(values) if values else 0
                        except (ValueError, TypeError):
                            pass
            
            aggregated.append(agg_row)
        
        new_state = TableState(data=aggregated)
        new_state.operation_chain = table_state.operation_chain + [
            f"AGGREGATE(group_by={group_by_column}, func={agg_func})"
        ]
        return new_state
    
    @staticmethod
    def sort_rows(
        table_state: TableState,
        column: str,
        order: str = "asc"
    ) -> TableState:
        """Sort table by a column (ascending or descending)."""
        reverse = order.lower() == "desc"
        try:
            # Try numeric sort first
            sorted_data = sorted(
                table_state.data,
                key=lambda row: float(row.get(column, 0)),
                reverse=reverse
            )
        except (ValueError, TypeError):
            # Fall back to string sort
            sorted_data = sorted(
                table_state.data,
                key=lambda row: str(row.get(column, "")),
                reverse=reverse
            )
        
        new_state = TableState(data=sorted_data)
        new_state.operation_chain = table_state.operation_chain + [
            f"SORT(column={column}, order={order})"
        ]
        return new_state
    
    @staticmethod
    def drop_column(
        table_state: TableState,
        column: str
    ) -> TableState:
        """Remove a column from the table."""
        new_data = []
        for row in table_state.data:
            new_row = {k: v for k, v in row.items() if k != column}
            new_data.append(new_row)
        
        new_state = TableState(data=new_data)
        new_state.operation_chain = table_state.operation_chain + [
            f"DROP_COLUMN({column})"
        ]
        return new_state


class ChainOfTableReasoner:
    """Main reasoning engine that orchestrates the three-stage pipeline."""
    
    def __init__(self, model: str = "gpt-4o-mini", max_iterations: int = 10):
        self.model = model
        self.max_iterations = max_iterations
        self.operations = OperationRegistry()
    
    def reason(self, initial_table: List[Dict], question: str) -> Tuple[str, List[str]]:
        """
        Execute the Chain-of-Table reasoning pipeline.
        Returns: (final_answer, operation_chain)
        """
        state = TableState(data=initial_table)
        
        for iteration in range(self.max_iterations):
            print(f"\n=== Iteration {iteration + 1} ===")
            print(f"Current table rows: {len(state.data)}, columns: {state.columns()}")
            
            # Stage 1: Select the next operation
            operation_name = self._stage1_select_operation(state, question)
            
            if operation_name == "DONE":
                print("Chain complete, moving to final answer extraction.")
                break
            
            # Stage 2: Generate arguments for the operation
            operation_args = self._stage2_generate_args(state, question, operation_name)
            
            if not operation_args:
                print(f"Failed to parse arguments for {operation_name}, terminating.")
                break
            
            print(f"Selected: {operation_name} with args {operation_args}")
            
            # Execute operation (deterministically)
            try:
                state = self._execute_operation(state, operation_name, operation_args)
                print(f"Table after operation: {len(state.data)} rows")
                
                # Check for termination conditions
                if len(state.data) == 0:
                    print("Table is empty, terminating.")
                    break
            except Exception as e:
                print(f"Error executing operation: {e}")
                break
        
        # Stage 3b: Final answer extraction
        final_answer = self._stage3_extract_answer(state, question)
        
        return final_answer, state.operation_chain
    
    def _stage1_select_operation(self, state: TableState, question: str) -> str:
        """Stage 1: LLM selects the next operation."""
        prompt = f"""You are reasoning over a table to answer a question using iterative transformations.

Current table ({len(state.data)} rows, columns: {', '.join(state.columns())}):
{state.to_string(max_rows=5)}

Original question: {question}

Operations performed so far:
{json.dumps(state.operation_chain, indent=2) if state.operation_chain else "None"}

What is the NEXT operation needed? Choose from:
1. ADD_COLUMN - Create a new column by extracting from an existing column
2. SELECT_ROWS - Keep only specific rows by index
3. FILTER - Keep rows matching a condition
4. AGGREGATE - Group and apply an aggregation function
5. SORT - Sort by a column
6. DROP_COLUMN - Remove an irrelevant column
7. DONE - The table is now ready to answer the question

Respond with ONLY the operation name (e.g., "ADD_COLUMN" or "DONE"). No explanation."""
        
        response = client.messages.create(
            model=self.model,
            max_tokens=50,
            messages=[{"role": "user", "content": prompt}]
        )
        
        operation_name = response.content[0].text.strip().upper()
        return operation_name
    
    def _stage2_generate_args(
        self,
        state: TableState,
        question: str,
        operation_name: str
    ) -> Dict[str, Any]:
        """Stage 2: LLM generates concrete arguments for the operation."""
        if operation_name == "ADD_COLUMN":
            prompt = f"""Given this table:
{state.to_string(max_rows=5)}

And the question: {question}

You need to execute: ADD_COLUMN

Provide:
1. new_column_name: The name of the column to create
2. source_column: Which column to extract from
3. extraction_rule: How to extract (e.g., "split by comma and take the right side")

Respond in this exact format:
new_column_name: [name]
source_column: [column]
extraction_rule: [rule]"""
        
        elif operation_name == "SELECT_ROWS":
            prompt = f"""Given this table:
{state.to_string(max_rows=5)}

And the question: {question}

You need to execute: SELECT_ROWS

Which row indices (0-indexed) are relevant? List them as a comma-separated list.
Respond with only the indices, e.g.: 0,1,2"""
        
        elif operation_name == "FILTER":
            prompt = f"""Given this table:
{state.to_string(max_rows=5)}

And the question: {question}

You need to execute: FILTER

Provide:
1. column: The column to filter on
2. condition: The condition (e.g., "rank <= 3")

Respond in this exact format:
column: [column_name]
condition: [condition]"""
        
        elif operation_name == "AGGREGATE":
            prompt = f"""Given this table:
{state.to_string(max_rows=5)}

And the question: {question}

You need to execute: AGGREGATE

Provide:
1. group_by_column: Which column to group by
2. agg_func: The aggregation function (count, sum, avg, max, or min)

Respond in this exact format:
group_by_column: [column]
agg_func: [function]"""
        
        elif operation_name == "SORT":
            prompt = f"""Given this table:
{state.to_string(max_rows=5)}

And the question: {question}

You need to execute: SORT

Provide:
1. column: Which column to sort by
2. order: asc or desc

Respond in this exact format:
column: [column_name]
order: [asc/desc]"""
        
        elif operation_name == "DROP_COLUMN":
            prompt = f"""Given this table:
{state.to_string(max_rows=5)}

And the question: {question}

You need to execute: DROP_COLUMN

Which column is irrelevant to answering the question?
Respond with only the column name."""
        
        else:
            return {}
        
        response = client.messages.create(
            model=self.model,
            max_tokens=200,
            messages=[{"role": "user", "content": prompt}]
        )
        
        response_text = response.content[0].text.strip()
        return self._parse_arguments(operation_name, response_text)
    
    def _parse_arguments(self, operation_name: str, response_text: str) -> Dict[str, Any]:
        """Parse LLM response into operation arguments."""
        args = {}
        
        try:
            if operation_name == "ADD_COLUMN":
                for line in response_text.split("\n"):
                    if "new_column_name:" in line.lower():
                        args["new_column_name"] = line.split(":", 1)[1].strip()
                    elif "source_column:" in line.lower():
                        args["source_column"] = line.split(":", 1)[1].strip()
                    elif "extraction_rule:" in line.lower():
                        args["extraction_rule"] = line.split(":", 1)[1].strip()
                return args if all(k in args for k in ["new_column_name", "source_column", "extraction_rule"]) else {}
            
            elif operation_name == "SELECT_ROWS":
                indices = [int(x.strip()) for x in response_text.split(",")]
                return {"indices": indices}
            
            elif operation_name == "FILTER":
                for line in response_text.split("\n"):
                    if "column:" in line.lower():
                        args["column"] = line.split(":", 1)[1].strip()
                    elif "condition:" in line.lower():
                        args["condition"] = line.split(":", 1)[1].strip()
                return args if all(k in args for k in ["column", "condition"]) else {}
            
            elif operation_name == "AGGREGATE":
                for line in response_text.split("\n"):
                    if "group_by_column:" in line.lower():
                        args["group_by_column"] = line.split(":", 1)[1].strip()
                    elif "agg_func:" in line.lower():
                        args["agg_func"] = line.split(":", 1)[1].strip()
                return args if all(k in args for k in ["group_by_column", "agg_func"]) else {}
            
            elif operation_name == "SORT":
                for line in response_text.split("\n"):
                    if "column:" in line.lower():
                        args["column"] = line.split(":", 1)[1].strip()
                    elif "order:" in line.lower():
                        args["order"] = line.split(":", 1)[1].strip()
                return args if all(k in args for k in ["column", "order"]) else {}
            
            elif operation_name == "DROP_COLUMN":
                return {"column": response_text.strip()}
        
        except (ValueError, IndexError):
            return {}
        
        return args
    
    def _execute_operation(
        self,
        state: TableState,
        operation_name: str,
        args: Dict[str, Any]
    ) -> TableState:
        """Execute a table operation and return the new state."""
        if operation_name == "ADD_COLUMN":
            return self.operations.add_column(
                state,
                args["new_column_name"],
                args["source_column"],
                args["extraction_rule"]
            )
        elif operation_name == "SELECT_ROWS":
            return self.operations.select_rows(state, args["indices"])
        elif operation_name == "FILTER":
            return self.operations.filter_rows(state, args["column"], args["condition"])
        elif operation_name == "AGGREGATE":
            return self.operations.aggregate(state, args["group_by_column"], args["agg_func"])
        elif operation_name == "SORT":
            return self.operations.sort_rows(state, args["column"], args.get("order", "asc"))
        elif operation_name == "DROP_COLUMN":
            return self.operations.drop_column(state