From 048f85d9186e16153c906f600fa3b122e05d3bc1 Mon Sep 17 00:00:00 2001
From: Vasu Bhog <vabhog@microsoft.com>
Date: Thu, 3 Dec 2020 19:37:22 -0600
Subject: [PATCH] Fix notebook unordered grid values after papermill execution
 (#13614)

* Fix unordered table

* check entire first row schema:

* SQL Notebooks should not get affected

* delete unused variable and edit comments

* refactor for efficient table ordering

* nit naming
---
 .../browser/outputs/gridOutput.component.ts   | 41 ++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts
index 12117fa9d8..37d816617d 100644
--- a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts
+++ b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts
@@ -40,7 +40,7 @@ import { ActionsOrientation } from 'vs/base/browser/ui/actionbar/actionbar';
 import { values } from 'vs/base/common/collections';
 import { URI } from 'vs/base/common/uri';
 import { assign } from 'vs/base/common/objects';
-
+import { equals } from 'vs/base/common/arrays';
 @Component({
 	selector: GridOutputComponent.SELECTOR,
 	template: `<div #output class="notebook-cellTable"></div>`
@@ -119,6 +119,7 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 		}
 		if (!this._table) {
 			let source = <IDataResource><any>this._bundleOptions.data[this.mimeType];
+			reorderGridData(source);
 			let state = new GridTableState(0, 0);
 			this._table = this.instantiationService.createInstance(DataResourceTable, source, this.cellModel, this.cellOutput, state);
 			let outputElement = <HTMLElement>this.output.nativeElement;
@@ -146,6 +147,44 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 	}
 }
 
+function reorderGridData(source: IDataResource): void {
+	// Get Column Names list from the data resource schema
+	const columnNames: string[] = source.schema.fields.map(field => field.name);
+	// Check to see if data source is ordered properly based on schema
+	// Papermill executed notebooks with KQL for instance will organize the
+	// row data in alphabetical order as a result the outputted grid will be
+	// unordered and not based on the data resource schema
+	if (source.data.length > 0) {
+		let rowKeys = Object.keys(source.data[0]);
+		if (!equals(columnNames, rowKeys)) {
+			// SQL notebooks do not use columnName as key (instead use indices)
+			// Indicies indicate the row is ordered properly
+			// We must check the data to know if it is in index form
+			let notIndexOrderKeys = false;
+			for (let index = 0; index < rowKeys.length - 1; index++) {
+				// Index form (all numbers, start at 0 and increase by 1)
+				let value = Number(rowKeys[index]);
+				if (isNaN(value) || value !== index) {
+					// break if key is not a number or in index form
+					notIndexOrderKeys = true;
+					break;
+				}
+			}
+			// Only reorder data that is not in index form
+			if (notIndexOrderKeys) {
+				source.data.forEach((row, index) => {
+					// Order each row based on the schema
+					let reorderedData = {};
+					for (let key of columnNames) {
+						reorderedData[key] = row[key];
+					}
+					source.data[index] = reorderedData;
+				});
+			}
+		}
+	}
+}
+
 class DataResourceTable extends GridTableBase<any> {
 
 	private _gridDataProvider: DataResourceDataProvider;