select_columns
Extract specific columns from a dataframe while removing all others. Validates column existence and reorders columns based on selection order, providing before/after column counts.
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 |
Implementation Reference
- The main handler function that executes the select_columns tool logic: validates columns exist, selects and copies them to the session dataframe, returns result with counts.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.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)