Basic helper for ES|QL's Apache Arrow output format (#2391)

(cherry picked from commit c3247d0)

Author: JoshMock

package.json

@@ -87,7 +87,8 @@
     "zx": "^7.2.2"
   },
   "dependencies": {
-    "@elastic/transport": "^8.8.1",
+    "@elastic/transport": "^8.9.0",
+    "@apache-arrow/esnext-cjs": "^17.0.0",
     "tslib": "^2.4.0"
   },
   "tap": {

src/helpers.ts

@@ -25,6 +25,7 @@ import assert from 'node:assert'
 import * as timersPromises from 'node:timers/promises'
 import { Readable } from 'node:stream'
 import { errors, TransportResult, TransportRequestOptions, TransportRequestOptionsWithMeta } from '@elastic/transport'
+import { Table, TypeMap, tableFromIPC } from '@apache-arrow/esnext-cjs'
 import Client from './client'
 import * as T from './api/types'
 
@@ -155,6 +156,7 @@ export interface EsqlResponse {
 
 export interface EsqlHelper {
   toRecords: <TDocument>() => Promise<EsqlToRecords<TDocument>>
+  toArrow: () => Promise<Table<TypeMap>>
 }
 
 export interface EsqlToRecords<TDocument> {
@@ -965,11 +967,6 @@ export default class Helpers {
    * @returns {object} EsqlHelper instance
    */
   esql (params: T.EsqlQueryRequest, reqOptions: TransportRequestOptions = {}): EsqlHelper {
-    if (this[kMetaHeader] !== null) {
-      reqOptions.headers = reqOptions.headers ?? {}
-      reqOptions.headers['x-elastic-client-meta'] = `${this[kMetaHeader] as string},h=qo`
-    }
-
     const client = this[kClient]
 
     function toRecords<TDocument> (response: EsqlResponse): TDocument[] {
@@ -985,17 +982,37 @@ export default class Helpers {
       })
     }
 
+    const metaHeader = this[kMetaHeader]
+
     const helper: EsqlHelper = {
       /**
        * Pivots ES|QL query results into an array of row objects, rather than the default format where each row is an array of values.
        */
       async toRecords<TDocument>(): Promise<EsqlToRecords<TDocument>> {
+        if (metaHeader !== null) {
+          reqOptions.headers = reqOptions.headers ?? {}
+          reqOptions.headers['x-elastic-client-meta'] = `${metaHeader as string},h=qo`
+        }
+
         params.format = 'json'
+        params.columnar = false
         // @ts-expect-error it's typed as ArrayBuffer but we know it will be JSON
         const response: EsqlResponse = await client.esql.query(params, reqOptions)
         const records: TDocument[] = toRecords(response)
         const { columns } = response
         return { records, columns }
+      },
+
+      async toArrow (): Promise<Table<TypeMap>> {
+        if (metaHeader !== null) {
+          reqOptions.headers = reqOptions.headers ?? {}
+          reqOptions.headers['x-elastic-client-meta'] = `${metaHeader as string},h=qa`
+        }
+
+        params.format = 'arrow'
+
+        const response = await client.esql.query(params, reqOptions)
+        return tableFromIPC(response)
       }
     }
 

test/unit/helpers/esql.test.ts

@@ -18,6 +18,7 @@
  */
 
 import { test } from 'tap'
+import { Table } from '@apache-arrow/esnext-cjs'
 import { connection } from '../../utils'
 import { Client } from '../../../'
 
@@ -109,5 +110,66 @@ test('ES|QL helper', t => {
 
     t.end()
   })
+
+  test('toArrow', t => {
+    t.test('Parses a binary response into an Arrow table', async t => {
+      const binaryContent = '/////zABAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAIAAAB8AAAABAAAAJ7///8UAAAARAAAAEQAAAAAAAoBRAAAAAEAAAAEAAAAjP///wgAAAAQAAAABAAAAGRhdGUAAAAADAAAAGVsYXN0aWM6dHlwZQAAAAAAAAAAgv///wAAAQAEAAAAZGF0ZQAAEgAYABQAEwASAAwAAAAIAAQAEgAAABQAAABMAAAAVAAAAAAAAwFUAAAAAQAAAAwAAAAIAAwACAAEAAgAAAAIAAAAEAAAAAYAAABkb3VibGUAAAwAAABlbGFzdGljOnR5cGUAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAGAAAAYW1vdW50AAAAAAAA/////7gAAAAUAAAAAAAAAAwAFgAOABUAEAAEAAwAAABgAAAAAAAAAAAABAAQAAAAAAMKABgADAAIAAQACgAAABQAAABYAAAABQAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAQAAAAAAAAAIAAAAAAAAACgAAAAAAAAAMAAAAAAAAAABAAAAAAAAADgAAAAAAAAAKAAAAAAAAAAAAAAAAgAAAAUAAAAAAAAAAAAAAAAAAAAFAAAAAAAAAAAAAAAAAAAAHwAAAAAAAAAAAACgmZkTQAAAAGBmZiBAAAAAAAAAL0AAAADAzMwjQAAAAMDMzCtAHwAAAAAAAADV6yywkgEAANWPBquSAQAA1TPgpZIBAADV17mgkgEAANV7k5uSAQAA/////wAAAAA='
+
+      const MockConnection = connection.buildMockConnection({
+        onRequest (_params) {
+          return {
+            body: Buffer.from(binaryContent, 'base64'),
+            statusCode: 200,
+            headers: {
+              'content-type': 'application/vnd.elasticsearch+arrow+stream'
+            }
+          }
+        }
+      })
+
+      const client = new Client({
+        node: 'http://localhost:9200',
+        Connection: MockConnection
+      })
+
+      const result = await client.helpers.esql({ query: 'FROM sample_data' }).toArrow()
+      t.ok(result instanceof Table)
+
+      const table = [...result]
+      t.same(table[0], [
+        ["amount", 4.900000095367432],
+        ["date", 1729532586965],
+      ])
+      t.end()
+    })
+
+    t.test('ESQL helper uses correct x-elastic-client-meta helper value', async t => {
+      const binaryContent = '/////zABAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAIAAAB8AAAABAAAAJ7///8UAAAARAAAAEQAAAAAAAoBRAAAAAEAAAAEAAAAjP///wgAAAAQAAAABAAAAGRhdGUAAAAADAAAAGVsYXN0aWM6dHlwZQAAAAAAAAAAgv///wAAAQAEAAAAZGF0ZQAAEgAYABQAEwASAAwAAAAIAAQAEgAAABQAAABMAAAAVAAAAAAAAwFUAAAAAQAAAAwAAAAIAAwACAAEAAgAAAAIAAAAEAAAAAYAAABkb3VibGUAAAwAAABlbGFzdGljOnR5cGUAAAAAAAAAAAAABgAIAAYABgAAAAAAAgAGAAAAYW1vdW50AAAAAAAA/////7gAAAAUAAAAAAAAAAwAFgAOABUAEAAEAAwAAABgAAAAAAAAAAAABAAQAAAAAAMKABgADAAIAAQACgAAABQAAABYAAAABQAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAQAAAAAAAAAIAAAAAAAAACgAAAAAAAAAMAAAAAAAAAABAAAAAAAAADgAAAAAAAAAKAAAAAAAAAAAAAAAAgAAAAUAAAAAAAAAAAAAAAAAAAAFAAAAAAAAAAAAAAAAAAAAHwAAAAAAAAAAAACgmZkTQAAAAGBmZiBAAAAAAAAAL0AAAADAzMwjQAAAAMDMzCtAHwAAAAAAAADV6yywkgEAANWPBquSAQAA1TPgpZIBAADV17mgkgEAANV7k5uSAQAA/////wAAAAA='
+
+      const MockConnection = connection.buildMockConnection({
+        onRequest (params) {
+          const header = params.headers?.['x-elastic-client-meta'] ?? ''
+          t.ok(header.includes('h=qa'), `Client meta header does not include ESQL helper value: ${header}`)
+          return {
+            body: Buffer.from(binaryContent, 'base64'),
+            statusCode: 200,
+            headers: {
+              'content-type': 'application/vnd.elasticsearch+arrow+stream'
+            }
+          }
+        }
+      })
+
+      const client = new Client({
+        node: 'http://localhost:9200',
+        Connection: MockConnection
+      })
+
+      await client.helpers.esql({ query: 'FROM sample_data' }).toArrow()
+      t.end()
+    })
+
+    t.end()
+  })
   t.end()
 })

tsconfig.json

@@ -1,6 +1,6 @@
 {
   "compilerOptions": {
-    "target": "es2019",
+    "target": "ES2019",
     "module": "commonjs",
     "moduleResolution": "node",
     "declaration": true,
@@ -21,7 +21,8 @@
     "importHelpers": true,
     "outDir": "lib",
     "lib": [
-      "esnext"
+      "ES2019",
+      "dom"
     ]
   },
   "formatCodeOptions": {