From 7cf5172b8a4901586ae23458d994c3094fee3521 Mon Sep 17 00:00:00 2001
From: Ryan McKinley <ryantxu@gmail.com>
Date: Tue, 30 Jun 2020 11:32:02 -0700
Subject: [PATCH] Arrow: when exporting arrow use the calculated names (#25871)

---
 .../src/dataframe/ArrowDataFrame.test.ts      | 26 ++++++++++++-
 .../src/dataframe/ArrowDataFrame.ts           | 38 +++++++++++++++++--
 2 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/packages/grafana-data/src/dataframe/ArrowDataFrame.test.ts b/packages/grafana-data/src/dataframe/ArrowDataFrame.test.ts
index 0910961d362..a77a4b6b0e6 100644
--- a/packages/grafana-data/src/dataframe/ArrowDataFrame.test.ts
+++ b/packages/grafana-data/src/dataframe/ArrowDataFrame.test.ts
@@ -3,8 +3,9 @@ import path from 'path';
 
 import { grafanaDataFrameToArrowTable, arrowTableToDataFrame } from './ArrowDataFrame';
 import { toDataFrameDTO, toDataFrame } from './processDataFrame';
-import { FieldType } from '../types';
+import { FieldType, DataFrame } from '../types';
 import { Table } from 'apache-arrow';
+import { ArrayVector } from '../vector';
 
 describe('Read/Write arrow Table to DataFrame', () => {
   test('should parse output with dataframe', () => {
@@ -58,4 +59,27 @@ describe('Read/Write arrow Table to DataFrame', () => {
     const frame = arrowTableToDataFrame(table);
     expect(toDataFrameDTO(frame)).toMatchSnapshot();
   });
+
+  test('Export arrow table names', () => {
+    const simple: DataFrame = {
+      name: 'hello',
+      fields: [
+        {
+          name: 'fname',
+          labels: {
+            a: 'AAA',
+            b: 'BBB',
+          },
+          config: {},
+          type: FieldType.number,
+          values: new ArrayVector([1, 2]),
+        },
+      ],
+      length: 2,
+    };
+    const t1 = grafanaDataFrameToArrowTable(simple);
+    const t2 = grafanaDataFrameToArrowTable(simple, true);
+    expect(t1.getColumnAt(0)?.name).toEqual('fname {a="AAA", b="BBB"}');
+    expect(t2.getColumnAt(0)?.name).toEqual('fname');
+  });
 });
diff --git a/packages/grafana-data/src/dataframe/ArrowDataFrame.ts b/packages/grafana-data/src/dataframe/ArrowDataFrame.ts
index 4aceeaf865e..05639b3995d 100644
--- a/packages/grafana-data/src/dataframe/ArrowDataFrame.ts
+++ b/packages/grafana-data/src/dataframe/ArrowDataFrame.ts
@@ -13,6 +13,7 @@ import {
   Bool,
   Column,
 } from 'apache-arrow';
+import { getFieldDisplayName } from '../field';
 
 export interface ArrowDataFrame extends DataFrame {
   table: Table;
@@ -117,16 +118,32 @@ function toArrowVector(field: Field): ArrowVector {
   return builder.finish().toVector();
 }
 
-export function grafanaDataFrameToArrowTable(data: DataFrame): Table {
+/**
+ * @param keepOriginalNames by default, the exported Table will get names that match the
+ * display within grafana.  This typically includes any labels defined in the metadata.
+ *
+ * When using this function to round-trip data, be sure to set `keepOriginalNames=true`
+ */
+export function grafanaDataFrameToArrowTable(data: DataFrame, keepOriginalNames?: boolean): Table {
   // Return the original table
   let table = (data as any).table;
-  if (table instanceof Table) {
-    return table as Table;
+  if (table instanceof Table && table.numCols === data.fields.length) {
+    if (!keepOriginalNames) {
+      table = updateArrowTableNames(table, data);
+    }
+    if (table) {
+      return table as Table;
+    }
   }
 
   table = Table.new(
     data.fields.map((field, index) => {
-      const column = Column.new(field.name, toArrowVector(field));
+      let name = field.name;
+      // when used directly as an arrow table the name should match the arrow schema
+      if (!keepOriginalNames) {
+        name = getFieldDisplayName(field, data);
+      }
+      const column = Column.new(name, toArrowVector(field));
       if (field.labels) {
         column.metadata.set('labels', JSON.stringify(field.labels));
       }
@@ -149,6 +166,19 @@ export function grafanaDataFrameToArrowTable(data: DataFrame): Table {
   return table;
 }
 
+function updateArrowTableNames(table: Table, frame: DataFrame): Table | undefined {
+  const cols: Column[] = [];
+  for (let i = 0; i < table.numCols; i++) {
+    const col = table.getColumnAt(i);
+    if (!col) {
+      return undefined;
+    }
+    const name = getFieldDisplayName(frame.fields[i], frame);
+    cols.push(Column.new(col.field.clone({ name: name }), ...col.chunks));
+  }
+  return Table.new(cols);
+}
+
 class NumberColumn extends FunctionalVector<number> {
   constructor(private col: Column) {
     super();