Data Explorer: Implement searchSchema API for DuckDB backend (#8894)

This brings the DuckDB backend up to feature parity with Python and R. I
can add sorting-by-data-type to the three backends in follow up PRs.

Author: wesm

extensions/positron-duckdb/src/extension.ts

@@ -8,6 +8,7 @@ import {
 	BackendState,
 	ColumnDisplayType,
 	ColumnFilter,
+	ColumnFilterType,
 	ColumnFrequencyTable,
 	ColumnFrequencyTableParams,
 	ColumnHistogram,
@@ -35,6 +36,7 @@ import {
 	FilterBetween,
 	FilterComparison,
 	FilterComparisonOp,
+	FilterMatchDataTypes,
 	FilterResult,
 	FilterSetMembership,
 	FilterTextSearch,
@@ -47,6 +49,9 @@ import {
 	ReturnColumnProfilesEvent,
 	RowFilter,
 	RowFilterType,
+	SearchSchemaParams,
+	SearchSchemaResult,
+	SearchSchemaSortOrder,
 	SetRowFiltersParams,
 	SetSortColumnsParams,
 	SupportStatus,
@@ -813,7 +818,120 @@ export class DuckDBTableView {
 					type_name: entry.column_type,
 					type_display
 				};
-			})
+			}),
+		};
+	}
+
+	async searchSchema(
+		params: SearchSchemaParams,
+	): RpcResponse<SearchSchemaResult> {
+		// Get all column indices
+		const allIndices: number[] = [];
+		for (let i = 0; i < this.fullSchema.length; i++) {
+			allIndices.push(i);
+		}
+
+		// Apply filters if any
+		let filteredIndices = allIndices;
+		if (params.filters && params.filters.length > 0) {
+			filteredIndices = allIndices.filter((index) => {
+				const entry = this.fullSchema[index];
+				const columnName = entry.column_name;
+				const columnType = entry.column_type;
+
+				// Get display type for this column
+				let displayType = SCHEMA_TYPE_MAPPING.get(columnType);
+				if (displayType === undefined) {
+					displayType = ColumnDisplayType.Unknown;
+				}
+				if (columnType.startsWith('DECIMAL')) {
+					displayType = ColumnDisplayType.Number;
+				}
+
+				// Apply each filter
+				return params.filters.every((filter) => {
+					switch (filter.filter_type) {
+						case ColumnFilterType.TextSearch: {
+							const textFilter =
+								filter.params as FilterTextSearch;
+							const searchTerm = textFilter.case_sensitive
+								? textFilter.term
+								: textFilter.term.toLowerCase();
+							const columnNameToMatch = textFilter.case_sensitive
+								? columnName
+								: columnName.toLowerCase();
+
+							switch (textFilter.search_type) {
+								case TextSearchType.Contains:
+									return columnNameToMatch.includes(
+										searchTerm,
+									);
+								case TextSearchType.NotContains:
+									return !columnNameToMatch.includes(
+										searchTerm,
+									);
+								case TextSearchType.StartsWith:
+									return columnNameToMatch.startsWith(
+										searchTerm,
+									);
+								case TextSearchType.EndsWith:
+									return columnNameToMatch.endsWith(
+										searchTerm,
+									);
+								case TextSearchType.RegexMatch:
+									try {
+										const regex = new RegExp(
+											textFilter.term,
+											textFilter.case_sensitive
+												? ''
+												: 'i',
+										);
+										return regex.test(columnName);
+									} catch {
+										return false;
+									}
+								default:
+									return false;
+							}
+						}
+						case ColumnFilterType.MatchDataTypes: {
+							const typeFilter =
+								filter.params as FilterMatchDataTypes;
+							return typeFilter.display_types.includes(
+								displayType,
+							);
+						}
+						default:
+							return false;
+					}
+				});
+			});
+		}
+
+		// Sort the filtered indices
+		switch (params.sort_order) {
+			case SearchSchemaSortOrder.Ascending:
+				filteredIndices.sort((a, b) => {
+					const nameA = this.fullSchema[a].column_name.toLowerCase();
+					const nameB = this.fullSchema[b].column_name.toLowerCase();
+					return nameA.localeCompare(nameB);
+				});
+				break;
+			case SearchSchemaSortOrder.Descending:
+				filteredIndices.sort((a, b) => {
+					const nameA = this.fullSchema[a].column_name.toLowerCase();
+					const nameB = this.fullSchema[b].column_name.toLowerCase();
+					return nameB.localeCompare(nameA);
+				});
+				break;
+			case SearchSchemaSortOrder.Original:
+			default:
+				// Keep original order
+				break;
+		}
+
+		return {
+			matches: filteredIndices,
 		};
 	}
 
@@ -1117,8 +1235,17 @@ END`;
 					]
 				},
 				search_schema: {
-					support_status: SupportStatus.Unsupported,
-					supported_types: []
+					support_status: SupportStatus.Supported,
+					supported_types: [
+						{
+							column_filter_type: ColumnFilterType.TextSearch,
+							support_status: SupportStatus.Supported,
+						},
+						{
+							column_filter_type: ColumnFilterType.MatchDataTypes,
+							support_status: SupportStatus.Supported,
+						}
+					],
 				},
 				set_column_filters: {
 					support_status: SupportStatus.Unsupported,
@@ -1272,11 +1399,12 @@ END`;
 								other_count: 0
 							};
 							break;
-						case ColumnProfileType.SummaryStats:
+						case ColumnProfileType.SummaryStats: {
 							// Create null summary stats appropriate for the column type
 							const columnSchema = this.fullSchema[request.column_index];
 							result.summary_stats = this.createEmptySummaryStats(columnSchema);
 							break;
+						}
 					}
 				}
 				return result;
@@ -1664,8 +1792,9 @@ export class DataExplorerRpcHandler implements vscode.Disposable {
 				return table.setRowFilters(rpc.params as SetRowFiltersParams);
 			case DataExplorerBackendRequest.SetSortColumns:
 				return table.setSortColumns(rpc.params as SetSortColumnsParams);
-			case DataExplorerBackendRequest.SetColumnFilters:
 			case DataExplorerBackendRequest.SearchSchema:
+				return table.searchSchema(rpc.params as SearchSchemaParams);
+			case DataExplorerBackendRequest.SetColumnFilters:
 				return `${rpc.method} not yet implemented`;
 			default:
 				return `unrecognized data explorer method: ${rpc.method} `;

extensions/positron-duckdb/src/interfaces.ts

@@ -16,6 +16,7 @@ export interface DataExplorerRpc {
 	uri?: string;
 	params: OpenDatasetParams |
 	GetSchemaParams |
+	SearchSchemaParams |
 	GetDataValuesParams |
 	GetRowLabelsParams |
 	GetColumnProfilesParams |
@@ -73,14 +74,10 @@ export interface OpenDatasetResult {
  */
 export interface SearchSchemaResult {
 	/**
-	 * A schema containing matching columns up to the max_results limit
+	 * The column indices of the matching column indices in the indicated
+	 * sort order
 	 */
-	matches: TableSchema;
-
-	/**
-	 * The total number of columns matching the filter
-	 */
-	total_num_matches: number;
+	matches: Array<number>;
 
 }
 
@@ -100,6 +97,28 @@ export interface ExportedData {
 
 }
 
+/**
+ * Code snippet for the data view
+ */
+export interface ConvertedCode {
+	/**
+	 * Lines of code that implement filters and sort keys
+	 */
+	converted_code: Array<string>;
+
+}
+
+/**
+ * Syntax to use for code conversion
+ */
+export interface CodeSyntaxName {
+	/**
+	 * The name of the code syntax, eg, pandas, polars, dplyr, etc.
+	 */
+	code_syntax_name: string;
+
+}
+
 /**
  * The result of applying filters to a table
  */
@@ -815,7 +834,7 @@ export interface ColumnFrequencyTable {
 	/**
 	 * The formatted top values
 	 */
-	values: Array<string>;
+	values: Array<ColumnValue>;
 
 	/**
 	 * Counts of top values
@@ -1006,21 +1025,19 @@ export interface SetSortColumnsFeatures {
 }
 
 /**
- * Feature flags for 'convert_to_code' RPC
+ * Feature flags for convert to code RPC
  */
 export interface ConvertToCodeFeatures {
 	/**
 	 * The support status for this RPC method
-	 * */
+	 */
 	support_status: SupportStatus;
+
 	/**
-	 * The supported code syntax names
+	 * The syntaxes for converted code
 	 */
-	supported_code_syntaxes?: Array<CodeSyntaxName>;
-}
+	code_syntaxes?: Array<CodeSyntaxName>;
 
-export interface CodeSyntaxName {
-	code_syntax_name: string;
 }
 
 /**
@@ -1143,6 +1160,15 @@ export type Selection = DataSelectionSingleCell | DataSelectionCellRange | DataS
 /// Union of selection specifications for array_selection
 export type ArraySelection = DataSelectionRange | DataSelectionIndices;
 
+/**
+ * Possible values for SortOrder in SearchSchema
+ */
+export enum SearchSchemaSortOrder {
+	Original = 'original',
+	Ascending = 'ascending',
+	Descending = 'descending'
+}
+
 /**
  * Possible values for ColumnDisplayType
  */
@@ -1294,20 +1320,15 @@ export interface GetSchemaParams {
  */
 export interface SearchSchemaParams {
 	/**
-	 * Column filters to apply when searching
+	 * Column filters to apply when searching, can be empty
 	 */
 	filters: Array<ColumnFilter>;
 
 	/**
-	 * Index (starting from zero) of first result to fetch (for paging)
-	 */
-	start_index: number;
-
-	/**
-	 * Maximum number of resulting column schemas to fetch from the start
-	 * index
+	 * How to sort results: original in-schema order, alphabetical ascending
+	 * or descending
 	 */
-	max_results: number;
+	sort_order: SearchSchemaSortOrder;
 }
 
 /**
@@ -1355,6 +1376,31 @@ export interface ExportDataSelectionParams {
 	format: ExportFormat;
 }
 
+/**
+ * Parameters for the ConvertToCode method.
+ */
+export interface ConvertToCodeParams {
+	/**
+	 * Zero or more column filters to apply
+	 */
+	column_filters: Array<ColumnFilter>;
+
+	/**
+	 * Zero or more row filters to apply
+	 */
+	row_filters: Array<RowFilter>;
+
+	/**
+	 * Zero or more sort keys to apply
+	 */
+	sort_keys: Array<ColumnSortKey>;
+
+	/**
+	 * The code syntax to use for conversion
+	 */
+	code_syntax_name: CodeSyntaxName;
+}
+
 /**
  * Parameters for the SetColumnFilters method.
  */
@@ -1471,6 +1517,8 @@ export enum DataExplorerBackendRequest {
 	GetDataValues = 'get_data_values',
 	GetRowLabels = 'get_row_labels',
 	ExportDataSelection = 'export_data_selection',
+	ConvertToCode = 'convert_to_code',
+	SuggestCodeSyntax = 'suggest_code_syntax',
 	SetColumnFilters = 'set_column_filters',
 	SetRowFilters = 'set_row_filters',
 	SetSortColumns = 'set_sort_columns',