Skip to content

Commit

Permalink
Return table shape from data explorer get_state RPC (#2345)
Browse files Browse the repository at this point in the history
* Move table shape to get_state RPC

* Put table shape in get_state RPC, refactor

* Rename parameter again

* revert submodule

* Update Python submodule
  • Loading branch information
wesm authored Feb 28, 2024
1 parent 59b7f95 commit a23199f
Show file tree
Hide file tree
Showing 11 changed files with 133 additions and 95 deletions.
2 changes: 1 addition & 1 deletion extensions/positron-python
38 changes: 24 additions & 14 deletions positron/comms/data_explorer-backend-openrpc.json
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,7 @@
"name": "table_schema",
"description": "The schema for a table-like object",
"required": [
"columns",
"num_rows",
"total_num_columns"
"columns"
],
"properties": {
"columns": {
Expand All @@ -44,14 +42,6 @@
"items": {
"$ref": "#/components/schemas/column_schema"
}
},
"num_rows": {
"type": "integer",
"description": "Numbers of rows in the unfiltered dataset"
},
"total_num_columns": {
"type": "integer",
"description": "Total number of columns in the unfiltered dataset"
}
}
}
Expand Down Expand Up @@ -279,18 +269,38 @@
{
"name": "get_state",
"summary": "Get the state",
"description": "Request the current backend state (applied filters and sort columns)",
"description": "Request the current table state (applied filters and sort columns)",
"params": [],
"result": {
"schema": {
"type": "object",
"name": "backend_state",
"description": "The current backend state",
"name": "table_state",
"description": "The current backend table state",
"required": [
"table_shape",
"filters",
"sort_keys"
],
"properties": {
"table_shape": {
"type": "object",
"name": "table_shape",
"description": "Provides number of rows and columns in table",
"required": [
"num_rows",
"num_columns"
],
"properties": {
"num_rows": {
"type": "integer",
"description": "Numbers of rows in the unfiltered dataset"
},
"num_columns": {
"type": "integer",
"description": "Number of columns in the unfiltered dataset"
}
}
},
"filters": {
"type": "array",
"description": "The set of currently applied filters",
Expand Down
2 changes: 1 addition & 1 deletion positron/comms/generate-comms.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ const tsOutputDir = `${__dirname}/../../src/vs/workbench/services/languageRuntim
const rustOutputDir = `${__dirname}/../../../amalthea/crates/amalthea/src/comm`;

/// The directory to write the generated Python files to
const pythonOutputDir = `${__dirname}/../../extensions/positron-python/pythonFiles/positron`;
const pythonOutputDir = `${__dirname}/../../extensions/positron-python/pythonFiles/positron/positron_ipykernel`;

const year = new Date().getFullYear();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,14 @@ import { Emitter, Event } from 'vs/base/common/event';
import { Disposable } from 'vs/base/common/lifecycle';
import { generateUuid } from 'vs/base/common/uuid';
import { IRuntimeClientInstance } from 'vs/workbench/services/languageRuntime/common/languageRuntimeClientInstance';
import { ColumnSortKey, PositronDataExplorerComm, SchemaUpdateEvent, TableData, TableSchema } from 'vs/workbench/services/languageRuntime/common/positronDataExplorerComm';
import {
ColumnSortKey,
PositronDataExplorerComm,
SchemaUpdateEvent,
TableData,
TableSchema,
TableState
} from 'vs/workbench/services/languageRuntime/common/positronDataExplorerComm';

/**
* A data explorer client instance.
Expand Down Expand Up @@ -81,6 +88,14 @@ export class DataExplorerClientInstance extends Disposable {
return this._positronDataExplorerComm.getSchema(startIndex, numColumns);
}

/**
* Get the current active state of the data explorer backend.
* @returns A promose that resolves to the current table state.
*/
async getState(): Promise<TableState> {
return this._positronDataExplorerComm.getState();
}

/**
* Get a rectangle of data values.
* @param rowStartIndex The first row to fetch (inclusive).
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,16 +19,6 @@ export interface TableSchema {
*/
columns: Array<ColumnSchema>;

/**
* Numbers of rows in the unfiltered dataset
*/
num_rows: number;

/**
* Total number of columns in the unfiltered dataset
*/
total_num_columns: number;

}

/**
Expand Down Expand Up @@ -126,9 +116,14 @@ export interface FreqtableCounts {
}

/**
* The current backend state
* The current backend table state
*/
export interface BackendState {
export interface TableState {
/**
* Provides number of rows and columns in table
*/
table_shape: TableShape;

/**
* The set of currently applied filters
*/
Expand All @@ -141,6 +136,22 @@ export interface BackendState {

}

/**
* Provides number of rows and columns in table
*/
export interface TableShape {
/**
* Numbers of rows in the unfiltered dataset
*/
num_rows: number;

/**
* Number of columns in the unfiltered dataset
*/
num_columns: number;

}

/**
* Schema for a column in a table
*/
Expand Down Expand Up @@ -447,12 +458,12 @@ export class PositronDataExplorerComm extends PositronBaseComm {
/**
* Get the state
*
* Request the current backend state (applied filters and sort columns)
* Request the current table state (applied filters and sort columns)
*
*
* @returns The current backend state
* @returns The current backend table state
*/
getState(): Promise<BackendState> {
getState(): Promise<TableState> {
return super.performRpc('get_state', [], []);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,12 @@ export class TableDataDataGridInstance extends DataGridInstance {
*/
private readonly _dataExplorerClientInstance: DataExplorerClientInstance;

private tableShape?: [number, number];

private _dataCache?: TableDataCache;
private _lastFetchedData?: FetchedData;

private _schemaCache: TableSchemaCache;
private _schemaCache?: TableSchemaCache;
private _lastFetchedSchema?: FetchedSchema;

//#endregion Private Properties
Expand Down Expand Up @@ -72,13 +74,6 @@ export class TableDataDataGridInstance extends DataGridInstance {
cellBorder: true
});

this._schemaCache = new TableSchemaCache(
async (req: SchemaFetchRange) => {
return this._dataExplorerClientInstance.getSchema(req.startIndex,
req.endIndex - req.startIndex);
}
);

// Set the data explorer client instance.
this._dataExplorerClientInstance = dataExplorerClientInstance;

Expand All @@ -92,12 +87,16 @@ export class TableDataDataGridInstance extends DataGridInstance {
this._firstColumnIndex = 0;
this._firstRowIndex = 0;

// Resets data schema, fetches initial schema and data
// Resets data schema, fetches table shape, initial schema, and data
this.initialize();
});

this._dataExplorerClientInstance.onDidDataUpdate(async (_evt) => {
this._lastFetchedData = undefined;

const state = await this._dataExplorerClientInstance.getState();
this.tableShape = [state.table_shape.num_rows, state.table_shape.num_columns];

this._dataCache?.clear();
this.fetchData();
});
Expand All @@ -109,42 +108,47 @@ export class TableDataDataGridInstance extends DataGridInstance {
* Gets the number of columns.
*/
get columns() {
return this._lastFetchedSchema ? this._lastFetchedSchema.schema.total_num_columns : 0;
return this.tableShape ? this.tableShape[1] : 0;
}

/**
* Gets the number of rows.
*/
get rows() {
return this._lastFetchedSchema ? this._lastFetchedSchema.schema.num_rows : 0;
return this.tableShape ? this.tableShape[0] : 0;
}

/**
*
*/
initialize() {
this._schemaCache.clear();
this._schemaCache.initialize().then(async (_) => {
this._lastFetchedSchema = await this._schemaCache?.fetch({ startIndex: 0, endIndex: 1000 });

this._dataCache = new TableDataCache(
this._schemaCache?.tableShape!,
async (req: DataFetchRange) => {
// Build the column indices to fetch.
const columnIndices: number[] = [];
for (let i = req.columnStartIndex; i < req.columnEndIndex; i++) {
columnIndices.push(i);
}
return this._dataExplorerClientInstance.getDataValues(
req.rowStartIndex,
req.rowEndIndex - req.rowStartIndex,
columnIndices
);
});

// Fetch data.
this.fetchData();
});
async initialize() {
const state = await this._dataExplorerClientInstance.getState();
this.tableShape = [state.table_shape.num_rows, state.table_shape.num_columns];
this._schemaCache = new TableSchemaCache(
this.tableShape,
async (req: SchemaFetchRange) => {
return this._dataExplorerClientInstance.getSchema(req.startIndex,
req.endIndex - req.startIndex);
}
);
this._lastFetchedSchema = await this._schemaCache?.fetch({ startIndex: 0, endIndex: 1000 });
this._dataCache = new TableDataCache(
this.tableShape,
async (req: DataFetchRange) => {
// Build the column indices to fetch.
const columnIndices: number[] = [];
for (let i = req.columnStartIndex; i < req.columnEndIndex; i++) {
columnIndices.push(i);
}
return this._dataExplorerClientInstance.getDataValues(
req.rowStartIndex,
req.rowEndIndex - req.rowStartIndex,
columnIndices
);
});

// Fetch data.
this.fetchData();
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ export class TableSummaryDataGridInstance extends DataGridInstance {
* Gets the number of rows.
*/
get rows() {
return this._tableSchema ? this._tableSchema.total_num_columns : 0;
return this._tableSchema ? this._tableSchema.columns.length : 0;
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,10 @@ export class TableDataCache extends FetchCache<DataFetchRange, FetchedData> {
this._fetchFunc = fetchFunc;
}

updateShape(newShape: [number, number]) {
this._tableShape = newShape;
}

getRangeTotalSize(range: DataFetchRange): number {
return ((range.columnEndIndex - range.columnStartIndex) *
(range.rowEndIndex - range.rowStartIndex));
Expand Down Expand Up @@ -197,17 +201,16 @@ export class TableSchemaCache extends FetchCache<SchemaFetchRange, FetchedSchema
private readonly SCHEMA_WINDOW = 50;

private _fetchFunc;
public tableShape: [number, number];
private _tableShape: [number, number];

constructor(fetchFunc: SchemaFetchFunc, maxCacheSize: number = 10_000) {
constructor(tableShape: [number, number], fetchFunc: SchemaFetchFunc, maxCacheSize: number = 10_000) {
super(maxCacheSize);
this.tableShape = [0, 0];
this._tableShape = tableShape;
this._fetchFunc = fetchFunc;
}

async initialize() {
const init_schema = await this._fetchFunc({ startIndex: 0, endIndex: 0 });
this.tableShape = [init_schema.num_rows, init_schema.total_num_columns];
updateShape(newShape: [number, number]) {
this._tableShape = newShape;
}

getRangeTotalSize(range: SchemaFetchRange): number {
Expand All @@ -223,7 +226,7 @@ export class TableSchemaCache extends FetchCache<SchemaFetchRange, FetchedSchema
range = structuredClone(range);

range.startIndex = Math.max(0, range.startIndex - this.SCHEMA_WINDOW);
range.endIndex = Math.min(this.tableShape[1], range.endIndex + this.SCHEMA_WINDOW);
range.endIndex = Math.min(this._tableShape[1], range.endIndex + this.SCHEMA_WINDOW);

const schema = await this._fetchFunc(range);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,6 @@ export function getTableSchema(numRows: number = 100, numColumns: number = 10):
}
return {
columns: columns,
num_rows: numRows,
total_num_columns: numColumns
};
}

Expand All @@ -53,13 +51,13 @@ const exampleTypeData: Record<string, string[]> = {
* @param numRows
* @param columnIndices Columns to select by index. Can be sequetial, sparse, or random
*/
export function getExampleTableData(schema: TableSchema, rowStartIndex: number,
export function getExampleTableData(shape: [number, number], schema: TableSchema, rowStartIndex: number,
numRows: number, columnIndices: Array<number>): TableData {
const generatedColumns = [];

// Don't generate virtual data beyond the extent of the table, and if
// rowStartIndex is at or after end of the table, return nothing
numRows = Math.max(Math.min(numRows, schema.num_rows - rowStartIndex), 0);
numRows = Math.max(Math.min(numRows, shape[0] - rowStartIndex), 0);

for (const columnIndex of columnIndices) {
const exampleValues: string[] = exampleTypeData[schema.columns[columnIndex].type_name];
Expand Down
Loading

0 comments on commit a23199f

Please sign in to comment.