From b87fc148d9f6d11d3b84b0ccfae37f7e1daeb78f Mon Sep 17 00:00:00 2001 From: Josh Mock Date: Mon, 21 Oct 2024 12:25:18 -0500 Subject: [PATCH 1/5] Ensure toRecords does not get columnar data --- src/helpers.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/src/helpers.ts b/src/helpers.ts index 62040083a..6d03d8bf8 100644 --- a/src/helpers.ts +++ b/src/helpers.ts @@ -991,6 +991,7 @@ export default class Helpers { */ async toRecords(): Promise> { 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) From fd216e4cf95289d5100369258113a8f91956d30c Mon Sep 17 00:00:00 2001 From: Josh Mock Date: Tue, 22 Oct 2024 14:14:29 -0500 Subject: [PATCH 2/5] Basic Apache Arrow support --- package.json | 1 + src/helpers.ts | 9 +++++++++ test/unit/helpers/esql.test.ts | 36 ++++++++++++++++++++++++++++++++++ tsconfig.json | 5 +++-- 4 files changed, 49 insertions(+), 2 deletions(-) diff --git a/package.json b/package.json index 2b4b5820a..7edb64262 100644 --- a/package.json +++ b/package.json @@ -88,6 +88,7 @@ }, "dependencies": { "@elastic/transport": "^8.8.1", + "@apache-arrow/esnext-cjs": "^17.0.0", "tslib": "^2.4.0" }, "tap": { diff --git a/src/helpers.ts b/src/helpers.ts index 6d03d8bf8..897e9883b 100644 --- a/src/helpers.ts +++ b/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, 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: () => Promise> + toArrow: () => Promise> } export interface EsqlToRecords { @@ -997,6 +999,13 @@ export default class Helpers { const records: TDocument[] = toRecords(response) const { columns } = response return { records, columns } + }, + + async toArrow (): Promise> { + params.format = 'arrow' + + const response = await client.esql.query(params, reqOptions) + return tableFromIPC(response) } } diff --git a/test/unit/helpers/esql.test.ts b/test/unit/helpers/esql.test.ts index b029e1323..1ff0d96f7 100644 --- a/test/unit/helpers/esql.test.ts +++ b/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,40 @@ 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.end() + }) t.end() }) diff --git a/tsconfig.json b/tsconfig.json index e93828bd8..a7d7a1352 100644 --- a/tsconfig.json +++ b/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": { From a8cbc0856dceace73225415bc5f13d6c595fd578 Mon Sep 17 00:00:00 2001 From: Josh Mock Date: Tue, 22 Oct 2024 14:19:07 -0500 Subject: [PATCH 3/5] Upgrade transport to 8.9.0 --- package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package.json b/package.json index 7edb64262..87b7fb9bc 100644 --- a/package.json +++ b/package.json @@ -87,7 +87,7 @@ "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" }, From 8899803fb8dedf230ffc67a8eafcb06b64f93b70 Mon Sep 17 00:00:00 2001 From: Josh Mock Date: Tue, 22 Oct 2024 14:41:11 -0500 Subject: [PATCH 4/5] Improve type definition for toArrow --- src/helpers.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/helpers.ts b/src/helpers.ts index 897e9883b..34679016d 100644 --- a/src/helpers.ts +++ b/src/helpers.ts @@ -25,7 +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, tableFromIPC } from '@apache-arrow/esnext-cjs' +import { Table, TypeMap, tableFromIPC } from '@apache-arrow/esnext-cjs' import Client from './client' import * as T from './api/types' @@ -156,7 +156,7 @@ export interface EsqlResponse { export interface EsqlHelper { toRecords: () => Promise> - toArrow: () => Promise> + toArrow: () => Promise> } export interface EsqlToRecords { @@ -1001,7 +1001,7 @@ export default class Helpers { return { records, columns } }, - async toArrow (): Promise> { + async toArrow (): Promise> { params.format = 'arrow' const response = await client.esql.query(params, reqOptions) From 413a0fa7c88816e9a5b63dfcc3d0a7db2718bc67 Mon Sep 17 00:00:00 2001 From: Josh Mock Date: Tue, 22 Oct 2024 14:50:57 -0500 Subject: [PATCH 5/5] Track Arrow helper usage in meta header --- src/helpers.ts | 17 ++++++++++++----- test/unit/helpers/esql.test.ts | 26 ++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/src/helpers.ts b/src/helpers.ts index 34679016d..a54ee0964 100644 --- a/src/helpers.ts +++ b/src/helpers.ts @@ -967,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 (response: EsqlResponse): TDocument[] { @@ -987,11 +982,18 @@ 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(): Promise> { + 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 @@ -1002,6 +1004,11 @@ export default class Helpers { }, async toArrow (): Promise> { + 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) diff --git a/test/unit/helpers/esql.test.ts b/test/unit/helpers/esql.test.ts index 1ff0d96f7..3685b7c53 100644 --- a/test/unit/helpers/esql.test.ts +++ b/test/unit/helpers/esql.test.ts @@ -143,6 +143,32 @@ test('ES|QL helper', t => { 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()