select_columns
Extract specific columns from a dataframe while removing all others, validating column existence and reordering by selection order.
Instructions
Select specific columns from dataframe, removing all others.
Validates column existence and reorders by selection order. Returns selection details with before/after column counts.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| columns | Yes | List of column names to select and keep |
Output Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| success | No | Whether operation completed successfully | |
| columns_after | Yes | Number of columns after selection | |
| columns_before | Yes | Number of columns before selection | |
| selected_columns | Yes | List of selected column names |
Implementation Reference
- The main handler function for the 'select_columns' tool. It retrieves the current dataframe from the session, validates that the specified columns exist, creates a new dataframe with only the selected columns, updates the session, and returns a SelectColumnsResult with details on the selection.
async def select_columns( ctx: Annotated[Context, Field(description="FastMCP context for session access")], columns: Annotated[list[str], Field(description="List of column names to select and keep")], ) -> SelectColumnsResult: """Select specific columns from dataframe, removing all others. Validates column existence and reorders by selection order. Returns selection details with before/after column counts. """ # Get session_id from FastMCP context session_id = ctx.session_id session, df = get_session_data(session_id) # Validate columns exist missing_cols = [col for col in columns if col not in df.columns] if missing_cols: raise ColumnNotFoundError(missing_cols[0], df.columns.tolist()) # Track counts before modification columns_before = len(df.columns) session.df = df[columns].copy() # No longer recording operations (simplified MCP architecture) return SelectColumnsResult( selected_columns=columns, columns_before=columns_before, columns_after=len(columns), ) - Pydantic model defining the response schema for the 'select_columns' tool, including selected columns and before/after column counts.
class SelectColumnsResult(BaseToolResponse): """Result of selecting specific columns.""" model_config = ConfigDict(extra="forbid") selected_columns: list[str] = Field(description="List of selected column names") columns_before: int = Field(description="Number of columns before selection") columns_after: int = Field(description="Number of columns after selection") - src/databeak/servers/column_server.py:653-653 (registration)Registration of the 'select_columns' tool on the column_server FastMCP instance.
column_server.tool(name="select_columns")(select_columns)