From 4cd27a380d89314118c4833177edf01fba647585 Mon Sep 17 00:00:00 2001
From: Ryan McKinley <ryantxu@gmail.com>
Date: Mon, 14 Mar 2022 22:58:41 -0700
Subject: [PATCH] Geomap: applying multiple line builders should keep each
 segment (#46563)

---
 public/app/features/geo/format/utils.ts       | 36 ++++++++++++++++++-
 .../spatial/spatialTransformer.ts             | 11 ++++--
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/public/app/features/geo/format/utils.ts b/public/app/features/geo/format/utils.ts
index 37ea8097d29..e3d9149d70e 100644
--- a/public/app/features/geo/format/utils.ts
+++ b/public/app/features/geo/format/utils.ts
@@ -1,6 +1,6 @@
 import { ArrayVector, Field, FieldConfig, FieldType } from '@grafana/data';
 import { getCenterPoint } from 'app/features/transformers/spatial/utils';
-import { Geometry, LineString, Point } from 'ol/geom';
+import { Geometry, GeometryCollection, LineString, Point } from 'ol/geom';
 import { fromLonLat } from 'ol/proj';
 import { Gazetteer } from '../gazetteer/gazetteer';
 import { decodeGeohash } from './geohash';
@@ -50,6 +50,40 @@ export function getGeoFieldFromGazetteer(gaz: Gazetteer, field: Field<string>):
   };
 }
 
+export function createGeometryCollection(
+  src: Field<Geometry | undefined>,
+  dest: Field<Geometry | undefined>
+): Field<Geometry | undefined> {
+  const v0 = src.values.toArray();
+  const v1 = dest.values.toArray();
+  if (!v0 || !v1) {
+    throw 'missing src/dest';
+  }
+  if (v0.length !== v1.length) {
+    throw 'Source and destination field lengths do not match';
+  }
+
+  const count = src.values.length!;
+  const geo = new Array<Geometry | undefined>(count);
+  for (let i = 0; i < count; i++) {
+    const a = v0[i];
+    const b = v1[i];
+    if (a && b) {
+      geo[i] = new GeometryCollection([a, b]);
+    } else if (a) {
+      geo[i] = a;
+    } else if (b) {
+      geo[i] = b;
+    }
+  }
+  return {
+    name: 'Geometry',
+    type: FieldType.geo,
+    values: new ArrayVector(geo),
+    config: hiddenTooltipField,
+  };
+}
+
 export function createLineBetween(
   src: Field<Geometry | undefined>,
   dest: Field<Geometry | undefined>
diff --git a/public/app/features/transformers/spatial/spatialTransformer.ts b/public/app/features/transformers/spatial/spatialTransformer.ts
index b9ce98d34f4..03b2a4d2dc3 100644
--- a/public/app/features/transformers/spatial/spatialTransformer.ts
+++ b/public/app/features/transformers/spatial/spatialTransformer.ts
@@ -1,5 +1,5 @@
 import { ArrayVector, DataFrame, DataTransformerID, DataTransformerInfo, FieldType } from '@grafana/data';
-import { createLineBetween } from 'app/features/geo/format/utils';
+import { createGeometryCollection, createLineBetween } from 'app/features/geo/format/utils';
 import { getGeometryField, getLocationMatchers } from 'app/features/geo/utils/location';
 import { mergeMap, from } from 'rxjs';
 import { SpatialOperation, SpatialAction, SpatialTransformOptions } from './models.gen';
@@ -26,10 +26,17 @@ async function doSetGeometry(frames: DataFrame[], options: SpatialTransformOptio
       const src = getGeometryField(frame, location);
       const target = getGeometryField(frame, targetLocation);
       if (src.field && target.field) {
+        const fields = [...frame.fields];
         const line = createLineBetween(src.field, target.field);
+        const first = fields[0];
+        if (first.type === FieldType.geo && first !== src.field && first !== target.field) {
+          fields[0] = createGeometryCollection(first, line); //
+        } else {
+          fields.unshift(line);
+        }
         return {
           ...frame,
-          fields: [line, ...frame.fields],
+          fields,
         };
       }
       return frame;