From aab29f7641aae38788b1d711a753d5e50deda03e Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 23:03:38 +0000 Subject: [PATCH 1/3] feat: add rowMode option to return rows as arrays Rows keyed by column name lose data when columns share a name (e.g. SELECT carat AS a, color AS a). rowMode: 'array' returns each row as values in schema order on the Thrift JSON, Arrow, CloudFetch and kernel paths. The default object mode is unchanged but now warns once per result set when it detects duplicate column names. Fixes #118 Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + lib/DBSQLSession.ts | 23 +- lib/contracts/IDBSQLSession.ts | 14 + lib/kernel/KernelOperationBackend.ts | 16 +- lib/kernel/KernelSessionBackend.ts | 5 +- lib/result/ArrowResultConverter.ts | 67 ++- lib/result/JsonResultHandler.ts | 25 +- lib/result/utils.ts | 24 + lib/thrift-backend/ThriftOperationBackend.ts | 10 +- lib/thrift-backend/ThriftSessionBackend.ts | 11 +- tests/unit/result/rowMode.test.ts | 478 +++++++++++++++++++ 11 files changed, 655 insertions(+), 19 deletions(-) create mode 100644 tests/unit/result/rowMode.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 77d44b38..35917fd5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Add the `rowMode` option to `executeStatement()`: `rowMode: 'array'` returns each row as an array of values in column order, so columns with the same name (e.g. `SELECT carat AS a, color AS a`) are no longer lost; the default `'object'` rows are unchanged, and the driver now logs a warning when a result has duplicate column names in object mode. (#118) - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: (last error: )`. `` is `HTTP : ` (or `HTTP ` when the status text is empty) for HTTP failures, and `network error `, `network error ()` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471) - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470) - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted. diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts index 18588ac3..0240db4d 100644 --- a/lib/DBSQLSession.ts +++ b/lib/DBSQLSession.ts @@ -29,6 +29,8 @@ import ISessionBackend from './contracts/ISessionBackend'; import IOperationBackend from './contracts/IOperationBackend'; import { numberToInt64 as numberToInt64Impl } from './thrift-backend/ThriftSessionBackend'; import { safeEmit } from './telemetry/telemetryUtils'; +import { ResultMetadata } from './contracts/ResultMetadata'; +import { getSchemaColumns } from './result/utils'; // Explicitly promisify a callback-style `pipeline` because `node:stream/promises` is not available in Node 14 const pipeline = util.promisify(stream.pipeline); @@ -91,6 +93,9 @@ export default class DBSQLSession implements IDBSQLSession { * const operation = await session.executeStatement(query); */ public async executeStatement(statement: string, options: ExecuteStatementOptions = {}): Promise { + if (options.rowMode !== undefined && options.rowMode !== 'object' && options.rowMode !== 'array') { + throw new HiveDriverError(`Unsupported rowMode: ${options.rowMode}. Expected 'object' or 'array'.`); + } const opBackend = await this.runBackend(() => this.backend.executeStatement(statement, options)); const operation = this.wrapOperation(opBackend); @@ -107,13 +112,17 @@ export default class DBSQLSession implements IDBSQLSession { const allowedLocalPath = Array.isArray(options.stagingAllowedLocalPath) ? options.stagingAllowedLocalPath : [options.stagingAllowedLocalPath]; - return this.handleStagingOperation(operation, allowedLocalPath); + return this.handleStagingOperation(operation, allowedLocalPath, metadata); } } return operation; } - private async handleStagingOperation(operation: IOperation, allowedLocalPath: Array): Promise { + private async handleStagingOperation( + operation: IOperation, + allowedLocalPath: Array, + metadata: ResultMetadata, + ): Promise { type StagingResponse = { presignedUrl: string; localFile?: string; @@ -124,7 +133,15 @@ export default class DBSQLSession implements IDBSQLSession { if (rows.length !== 1) { throw new StagingError('Staging operation: expected only one row in result'); } - const row = rows[0] as StagingResponse; + let row = rows[0] as StagingResponse; + // `rowMode: 'array'` returns positional rows; key them by column name. + if (Array.isArray(row)) { + const values: Array = row; + row = getSchemaColumns(metadata.schema).reduce( + (result, column, i) => ({ ...result, [column.columnName]: values[i] }), + {} as StagingResponse, + ); + } if (row.localFile !== undefined) { let allowOperation = false; diff --git a/lib/contracts/IDBSQLSession.ts b/lib/contracts/IDBSQLSession.ts index e7eb31d8..12d7d6c2 100644 --- a/lib/contracts/IDBSQLSession.ts +++ b/lib/contracts/IDBSQLSession.ts @@ -4,6 +4,15 @@ import Status from '../dto/Status'; import InfoValue from '../dto/InfoValue'; import { DBSQLParameter, DBSQLParameterValue } from '../DBSQLParameter'; +/** + * Shape of the rows returned by a statement's fetch methods: + * - `'object'` (default): each row is an object keyed by column name. Columns + * that share a name (e.g. `SELECT a.id, b.id`) overwrite each other. + * - `'array'`: each row is an array of values in result-schema order, so every + * column is kept. Use `getResultMetadata()` / `getSchema()` for column names. + */ +export type RowMode = 'object' | 'array'; + export type ExecuteStatementOptions = { /** * The number of seconds after which the query will time out on the server. @@ -50,6 +59,11 @@ export type ExecuteStatementOptions = { * `query_tags` key). Ignored by the Thrift backend. */ statementConf?: Record; + /** + * Row shape returned by this statement's `fetchChunk` / `fetchAll` / iterators. + * Defaults to `'object'`; see {@link RowMode}. + */ + rowMode?: RowMode; }; export type TypeInfoRequest = { diff --git a/lib/kernel/KernelOperationBackend.ts b/lib/kernel/KernelOperationBackend.ts index 8173a0fe..8d6e3a65 100644 --- a/lib/kernel/KernelOperationBackend.ts +++ b/lib/kernel/KernelOperationBackend.ts @@ -43,6 +43,7 @@ import { OperationStatus, OperationState } from '../contracts/OperationStatus'; import { ResultMetadata, ResultFormat } from '../contracts/ResultMetadata'; import IClientContext from '../contracts/IClientContext'; import { LogLevel } from '../contracts/IDBSQLLogger'; +import { RowMode } from '../contracts/IDBSQLSession'; import Status from '../dto/Status'; import HiveDriverError from '../errors/HiveDriverError'; import OperationStateError, { OperationStateErrorCode } from '../errors/OperationStateError'; @@ -167,6 +168,8 @@ export interface KernelOperationBackendOptions { * handle exposes one, else a fresh UUIDv4. */ id?: string; + /** Row shape returned by `fetchChunk`; defaults to `'object'`. */ + rowMode?: RowMode; } export default class KernelOperationBackend implements IOperationBackend { @@ -202,6 +205,8 @@ export default class KernelOperationBackend implements IOperationBackend { private readonly lifecycle: KernelOperationLifecycleState = createLifecycleState(); + private readonly rowMode?: RowMode; + private resultSlicer?: ResultSlicer; private resultsProvider?: KernelResultsProvider; @@ -215,7 +220,14 @@ export default class KernelOperationBackend implements IOperationBackend { // already-terminal statement. Drives both fetch and result-metadata. private fetchHandlePromise?: Promise; - constructor({ asyncStatement, statement, cancellableExecution, context, id }: KernelOperationBackendOptions) { + constructor({ + asyncStatement, + statement, + cancellableExecution, + context, + id, + rowMode, + }: KernelOperationBackendOptions) { // Exactly one of the three handle kinds must be supplied. const providedCount = (asyncStatement !== undefined ? 1 : 0) + @@ -244,6 +256,7 @@ export default class KernelOperationBackend implements IOperationBackend { } : ((asyncStatement ?? statement) as KernelStatementHandle); this.context = context; + this.rowMode = rowMode; this._id = id ?? asyncStatement?.statementId ?? statement?.statementId ?? cancellableExecution?.statementId ?? uuidv4(); } @@ -789,6 +802,7 @@ export default class KernelOperationBackend implements IOperationBackend { const converter = new ArrowResultConverter(this.context, this.resultsProvider, metadata, { preserveBigNumericPrecision: this.context.getConfig().preserveBigNumericPrecision ?? false, disableRowMaterialization: this.context.getConfig().disableRowMaterialization ?? false, + rowMode: this.rowMode, }); this.resultSlicer = new ResultSlicer(this.context, converter); return this.resultSlicer; diff --git a/lib/kernel/KernelSessionBackend.ts b/lib/kernel/KernelSessionBackend.ts index 73fd5cb2..abe85c15 100644 --- a/lib/kernel/KernelSessionBackend.ts +++ b/lib/kernel/KernelSessionBackend.ts @@ -257,9 +257,9 @@ export default class KernelSessionBackend implements ISessionBackend { // Feature-detect the arm via a type guard: only the Running `AsyncStatement` // exposes `awaitResult`; the terminal `Statement` (Completed arm) does not. if (isKernelAsyncStatement(direct)) { - return new KernelOperationBackend({ asyncStatement: direct, context: this.context }); + return new KernelOperationBackend({ asyncStatement: direct, context: this.context, rowMode: options.rowMode }); } - return new KernelOperationBackend({ statement: direct, context: this.context }); + return new KernelOperationBackend({ statement: direct, context: this.context, rowMode: options.rowMode }); } // Async path: submit (`wait_timeout=0s`) returns a pending handle the @@ -279,6 +279,7 @@ export default class KernelSessionBackend implements ISessionBackend { return new KernelOperationBackend({ asyncStatement: asyncStatement!, context: this.context, + rowMode: options.rowMode, }); } diff --git a/lib/result/ArrowResultConverter.ts b/lib/result/ArrowResultConverter.ts index 5715d725..db296cb3 100644 --- a/lib/result/ArrowResultConverter.ts +++ b/lib/result/ArrowResultConverter.ts @@ -17,9 +17,10 @@ import { } from 'apache-arrow'; import { TTableSchema, TColumnDesc, TTypeId } from '../../thrift/TCLIService_types'; import IClientContext from '../contracts/IClientContext'; +import { RowMode } from '../contracts/IDBSQLSession'; import HiveDriverError from '../errors/HiveDriverError'; import IResultsProvider, { ResultsProviderFetchNextOptions } from './IResultsProvider'; -import { ArrowBatch, getSchemaColumns, convertThriftValue } from './utils'; +import { ArrowBatch, getSchemaColumns, convertThriftValue, warnOnDuplicateColumnNames } from './utils'; const { isArrowBigNumSymbol, bigNumToBigInt } = arrowUtils; @@ -208,6 +209,8 @@ export default class ArrowResultConverter implements IResultsProvider // (only the row count is meaningful). See `ConnectionOptions.disableRowMaterialization`. private readonly disableRowMaterialization: boolean; + private readonly rowMode: RowMode; + private recordBatchReader?: IterableIterator>; // Remaining rows in current Arrow batch (not the record batch!) @@ -232,13 +235,18 @@ export default class ArrowResultConverter implements IResultsProvider { preserveBigNumericPrecision = false, disableRowMaterialization = false, - }: { preserveBigNumericPrecision?: boolean; disableRowMaterialization?: boolean } = {}, + rowMode = 'object', + }: { preserveBigNumericPrecision?: boolean; disableRowMaterialization?: boolean; rowMode?: RowMode } = {}, ) { this.context = context; this.source = source; this.schema = getSchemaColumns(schema); this.preserveBigNumericPrecision = preserveBigNumericPrecision; this.disableRowMaterialization = disableRowMaterialization; + this.rowMode = rowMode; + if (rowMode === 'object' && !disableRowMaterialization) { + warnOnDuplicateColumnNames(context, this.schema); + } } public async hasMore() { @@ -268,15 +276,21 @@ export default class ArrowResultConverter implements IResultsProvider if (this.prefetchedRecordBatch) { // Consume a record batch fetched during previous call to `fetchNext` - const table = new Table(this.prefetchedRecordBatch); + const batch = this.prefetchedRecordBatch; this.prefetchedRecordBatch = undefined; // Get table rows, but not more than remaining count. When materialization // is disabled, skip `toArray()`/`getRows()` and return `null` placeholders // (filled, not sparse, so `fetchAll`'s `.flat()` keeps the count). - const batchRows = Math.min(table.numRows, this.remainingRows); - const result = this.disableRowMaterialization - ? new Array(batchRows).fill(null) - : this.getRows(table.schema, table.toArray().slice(0, this.remainingRows)); + const batchRows = Math.min(batch.numRows, this.remainingRows); + let result: Array; + if (this.disableRowMaterialization) { + result = new Array(batchRows).fill(null); + } else if (this.rowMode === 'array') { + result = this.getRowArrays(batch, batchRows); + } else { + const table = new Table(batch); + result = this.getRows(table.schema, table.toArray().slice(0, this.remainingRows)); + } // Reduce remaining rows count by a count of rows we just processed. // If the remaining count reached zero - we're done with current arrow @@ -345,6 +359,29 @@ export default class ArrowResultConverter implements IResultsProvider }); } + // Reads cells by column index, so same-named columns are all kept. Fields come from the + // batch data type: `RecordBatch.schema` merges same-named fields into one type. + private getRowArrays(batch: RecordBatch, rowCount: number): Array> { + const fields: Array = batch.data.type.children; + const vectors = this.schema.map((column, index) => batch.getChildAt(index)); + const rows: Array> = []; + for (let i = 0; i < rowCount; i += 1) { + const row: Array = []; + for (let j = 0; j < this.schema.length; j += 1) { + const field = fields[j]; + const value = this.convertArrowTypes( + vectors[j]?.get(i) ?? null, + field?.type, + field?.type.children || [], + field, + ); + row.push(this.convertThriftColumnValue(this.schema[j], value)); + } + rows.push(row); + } + return rows; + } + private convertArrowTypes( value: any, valueType: DataType | undefined, @@ -489,4 +526,20 @@ export default class ArrowResultConverter implements IResultsProvider return result; } + + // Per-cell equivalent of `convertThriftTypes`, used by array rows. Kept separate so + // the default object path stays as is. + private convertThriftColumnValue(column: TColumnDesc, value: any): any { + if (value === null) { + return null; + } + const typeDescriptor = column.typeDesc.types[0]?.primitiveEntry; + if ( + this.preserveBigNumericPrecision && + (typeDescriptor?.type === TTypeId.DECIMAL_TYPE || typeDescriptor?.type === TTypeId.BIGINT_TYPE) + ) { + return value; + } + return convertThriftValue(typeDescriptor, value); + } } diff --git a/lib/result/JsonResultHandler.ts b/lib/result/JsonResultHandler.ts index 2286b6dd..87f6b07d 100644 --- a/lib/result/JsonResultHandler.ts +++ b/lib/result/JsonResultHandler.ts @@ -1,7 +1,8 @@ import { TGetResultSetMetadataResp, TRowSet, TColumn, TColumnDesc } from '../../thrift/TCLIService_types'; import IClientContext from '../contracts/IClientContext'; +import { RowMode } from '../contracts/IDBSQLSession'; import IResultsProvider, { ResultsProviderFetchNextOptions } from './IResultsProvider'; -import { getSchemaColumns, convertThriftValue, getColumnValue } from './utils'; +import { getSchemaColumns, convertThriftValue, getColumnValue, warnOnDuplicateColumnNames } from './utils'; export default class JsonResultHandler implements IResultsProvider> { private readonly context: IClientContext; @@ -10,14 +11,21 @@ export default class JsonResultHandler implements IResultsProvider> { private readonly schema: Array; + private readonly rowMode: RowMode; + constructor( context: IClientContext, source: IResultsProvider, { schema }: TGetResultSetMetadataResp, + { rowMode = 'object' }: { rowMode?: RowMode } = {}, ) { this.context = context; this.source = source; this.schema = getSchemaColumns(schema); + this.rowMode = rowMode; + if (rowMode === 'object') { + warnOnDuplicateColumnNames(context, this.schema); + } } public async hasMore() { @@ -35,7 +43,20 @@ export default class JsonResultHandler implements IResultsProvider> { } const columns = data.columns || []; - return this.getRows(columns, this.schema); + return this.rowMode === 'array' ? this.getRowArrays(columns, this.schema) : this.getRows(columns, this.schema); + } + + private getRowArrays(columns: Array, descriptors: Array): Array> { + const rows: Array> = []; + descriptors.forEach((descriptor, columnIndex) => { + this.getSchemaValues(descriptor, columns[descriptor.position - 1]).forEach((value, i) => { + if (!rows[i]) { + rows[i] = []; + } + rows[i][columnIndex] = value; + }); + }); + return rows; } private getRows(columns: Array, descriptors: Array): Array { diff --git a/lib/result/utils.ts b/lib/result/utils.ts index 4bbdf418..cac927af 100644 --- a/lib/result/utils.ts +++ b/lib/result/utils.ts @@ -18,6 +18,8 @@ import { } from 'apache-arrow'; import { TTableSchema, TColumnDesc, TPrimitiveTypeEntry, TTypeId, TColumn } from '../../thrift/TCLIService_types'; import HiveDriverError from '../errors/HiveDriverError'; +import IClientContext from '../contracts/IClientContext'; +import { LogLevel } from '../contracts/IDBSQLLogger'; export interface ArrowBatch { batches: Array; @@ -32,6 +34,28 @@ export function getSchemaColumns(schema?: TTableSchema): Array { return [...schema.columns].sort((c1, c2) => c1.position - c2.position); } +// Object rows are keyed by column name, so same-named columns overwrite each other. +export function warnOnDuplicateColumnNames(context: IClientContext, columns: Array): void { + const seen = new Set(); + const duplicates = new Set(); + for (const { columnName } of columns) { + if (seen.has(columnName)) { + duplicates.add(columnName); + } + seen.add(columnName); + } + if (duplicates.size > 0) { + const names = [...duplicates].map((name) => JSON.stringify(name)).join(', '); + context + .getLogger() + .log( + LogLevel.warn, + `Result has duplicate column names (${names}); object rows keep only the last column with each name. ` + + "Alias the columns or pass `rowMode: 'array'` to executeStatement to keep every column.", + ); + } +} + function isString(value: unknown): value is string { return typeof value === 'string' || value instanceof String; } diff --git a/lib/thrift-backend/ThriftOperationBackend.ts b/lib/thrift-backend/ThriftOperationBackend.ts index dd5ac327..86ce6163 100644 --- a/lib/thrift-backend/ThriftOperationBackend.ts +++ b/lib/thrift-backend/ThriftOperationBackend.ts @@ -12,6 +12,7 @@ import { } from '../../thrift/TCLIService_types'; import IOperationBackend, { IOperationBackendWaitOptions } from '../contracts/IOperationBackend'; import IClientContext from '../contracts/IClientContext'; +import { RowMode } from '../contracts/IDBSQLSession'; import { OperationStatus, OperationState } from '../contracts/OperationStatus'; import { ResultMetadata, ResultFormat } from '../contracts/ResultMetadata'; import Status from '../dto/Status'; @@ -35,6 +36,7 @@ interface ThriftOperationBackendOptions { context: IClientContext; functionCatalog?: string | null; emptyResultOnNotFound?: TTableSchema; + rowMode?: RowMode; } async function delay(ms?: number): Promise { @@ -85,6 +87,8 @@ export default class ThriftOperationBackend implements IOperationBackend { private readonly emptyResultOnNotFound?: TTableSchema; + private readonly rowMode?: RowMode; + private readonly operationHandle: TOperationHandle; private readonly _data: RowSetProvider; @@ -107,11 +111,13 @@ export default class ThriftOperationBackend implements IOperationBackend { context, functionCatalog, emptyResultOnNotFound, + rowMode, }: ThriftOperationBackendOptions) { this.operationHandle = handle; this.context = context; this.functionCatalog = functionCatalog; this.emptyResultOnNotFound = emptyResultOnNotFound; + this.rowMode = rowMode; this.metadata = directResults?.resultSetMetadata; const useOnlyPrefetchedResults = Boolean(directResults?.closeOperation); @@ -351,7 +357,7 @@ export default class ThriftOperationBackend implements IOperationBackend { switch (resultFormat) { case TSparkRowSetType.COLUMN_BASED_SET: - resultSource = new JsonResultHandler(this.context, this._data, metadata); + resultSource = new JsonResultHandler(this.context, this._data, metadata, { rowMode: this.rowMode }); break; case TSparkRowSetType.ARROW_BASED_SET: resultSource = new ArrowResultConverter( @@ -361,6 +367,7 @@ export default class ThriftOperationBackend implements IOperationBackend { { preserveBigNumericPrecision: this.context.getConfig().preserveBigNumericPrecision ?? false, disableRowMaterialization: this.context.getConfig().disableRowMaterialization ?? false, + rowMode: this.rowMode, }, ); break; @@ -372,6 +379,7 @@ export default class ThriftOperationBackend implements IOperationBackend { { preserveBigNumericPrecision: this.context.getConfig().preserveBigNumericPrecision ?? false, disableRowMaterialization: this.context.getConfig().disableRowMaterialization ?? false, + rowMode: this.rowMode, }, ); break; diff --git a/lib/thrift-backend/ThriftSessionBackend.ts b/lib/thrift-backend/ThriftSessionBackend.ts index 8ccb1b96..49ec539a 100644 --- a/lib/thrift-backend/ThriftSessionBackend.ts +++ b/lib/thrift-backend/ThriftSessionBackend.ts @@ -26,6 +26,7 @@ import { FunctionsRequest, PrimaryKeysRequest, CrossReferenceRequest, + RowMode, } from '../contracts/IDBSQLSession'; import Status from '../dto/Status'; import InfoValue from '../dto/InfoValue'; @@ -224,7 +225,7 @@ export default class ThriftSessionBackend implements ISessionBackend { } const response = await driver.executeStatement(request); - return this.createOperationBackend(response); + return this.createOperationBackend(response, { rowMode: options.rowMode }); } public async getTypeInfo(request: TypeInfoRequest): Promise { @@ -308,7 +309,7 @@ export default class ThriftSessionBackend implements ISessionBackend { runAsync: this.getRunAsyncForMetadataOperations(), ...getDirectResultsOptions(request.maxRows, this.context.getConfig()), }); - return this.createOperationBackend(response, functionCatalog); + return this.createOperationBackend(response, { functionCatalog }); } public async getPrimaryKeys(request: PrimaryKeysRequest): Promise { @@ -373,7 +374,10 @@ export default class ThriftSessionBackend implements ISessionBackend { return new Status(response.status); } - private createOperationBackend(response: OperationResponseShape, functionCatalog?: string | null): IOperationBackend { + private createOperationBackend( + response: OperationResponseShape, + { functionCatalog, rowMode }: { functionCatalog?: string | null; rowMode?: RowMode } = {}, + ): IOperationBackend { Status.assert(response.status); const handle = definedOrError(response.operationHandle); return new ThriftOperationBackend({ @@ -381,6 +385,7 @@ export default class ThriftSessionBackend implements ISessionBackend { directResults: response.directResults, context: this.context, functionCatalog, + rowMode, }); } } diff --git a/tests/unit/result/rowMode.test.ts b/tests/unit/result/rowMode.test.ts new file mode 100644 index 00000000..3a0c5da7 --- /dev/null +++ b/tests/unit/result/rowMode.test.ts @@ -0,0 +1,478 @@ +import { expect } from 'chai'; +import sinon from 'sinon'; +import Int64 from 'node-int64'; +import { + Decimal, + Field, + Float64, + Int32, + Int64 as ArrowInt64, + makeData, + RecordBatch, + RecordBatchStreamWriter, + Schema, + Struct, + TimestampMillisecond, + Utf8, + vectorFromArray, +} from 'apache-arrow'; +import { + TColumnDesc, + TRowSet, + TSparkRowSetType, + TStatusCode, + TTableSchema, + TTypeId, +} from '../../../thrift/TCLIService_types'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; +import ArrowResultConverter from '../../../lib/result/ArrowResultConverter'; +import ArrowResultHandler from '../../../lib/result/ArrowResultHandler'; +import JsonResultHandler from '../../../lib/result/JsonResultHandler'; +import { ArrowBatch, getSchemaColumns } from '../../../lib/result/utils'; +import KernelOperationBackend from '../../../lib/kernel/KernelOperationBackend'; +import HiveDriverError from '../../../lib/errors/HiveDriverError'; +import ResultsProviderStub from '../.stubs/ResultsProviderStub'; +import ClientContextStub from '../.stubs/ClientContextStub'; +import { createSessionForTest } from '../.stubs/createSessionForTest'; +import { fixArrowResult } from '../../fixtures/compatibility'; +import * as fixtureColumn from '../../fixtures/compatibility/column'; +import * as fixtureArrow from '../../fixtures/compatibility/arrow'; +import * as fixtureArrowNT from '../../fixtures/compatibility/arrow_native_types'; + +function column(columnName: string, type: TTypeId, position: number): TColumnDesc { + return { columnName, typeDesc: { types: [{ primitiveEntry: { type } }] }, position }; +} + +// `SELECT carat AS a, color AS a` from the issue report. +const duplicateThriftSchema: TTableSchema = { + columns: [column('a', TTypeId.DOUBLE_TYPE, 1), column('a', TTypeId.STRING_TYPE, 2)], +}; + +const duplicateArrowSchema = new Schema([new Field('a', new Float64(), true), new Field('a', new Utf8(), true)]); + +function arrowIpc(schema: Schema, columns: Array>): Buffer { + // `new RecordBatch(schema, ...)` merges same-named fields, so build it with unique + // names and attach the real schema before serializing. + const uniqueFields = schema.fields.map((field, i) => new Field(`c${i}`, field.type, field.nullable, field.metadata)); + const children = schema.fields.map((field, i) => vectorFromArray(columns[i], field.type).data[0]); + const struct = makeData({ + type: new Struct(uniqueFields), + children, + length: children[0]?.length ?? 0, + nullCount: 0, + }); + const batch = new RecordBatch(new Schema(uniqueFields), struct); + Object.assign(batch, { schema }); + const writer = new RecordBatchStreamWriter(); + writer.write(batch); + writer.finish(); + return Buffer.from(writer.toUint8Array(true)); +} + +function duplicateRowSet(): TRowSet { + return { + startRowOffset: new Int64(0), + rows: [], + columns: [ + { doubleVal: { values: [0.23, 0.21], nulls: Buffer.from([]) } }, + { stringVal: { values: ['E', 'F'], nulls: Buffer.from([0b10]) } }, + ], + }; +} + +function toArrays(rows: Array>, schema: TTableSchema) { + const columns = getSchemaColumns(schema); + return rows.map((row) => columns.map((c) => row[c.columnName])); +} + +function warnings(context: ClientContextStub, spy: sinon.SinonSpy) { + return spy.getCalls().filter((call) => call.args[0] === LogLevel.warn && /duplicate column/i.test(call.args[1])); +} + +describe('rowMode', () => { + describe('JsonResultHandler', () => { + it('keeps every column of duplicate-named columns in array mode', async () => { + const handler = new JsonResultHandler( + new ClientContextStub(), + new ResultsProviderStub([duplicateRowSet()], undefined), + { schema: duplicateThriftSchema, status: { statusCode: TStatusCode.SUCCESS_STATUS } }, + { rowMode: 'array' }, + ); + expect(await handler.fetchNext({ limit: 100 })).to.deep.equal([ + [0.23, 'E'], + [0.21, null], + ]); + }); + + it('orders array values by column position', async () => { + const schema: TTableSchema = { + columns: [column('second', TTypeId.STRING_TYPE, 2), column('first', TTypeId.INT_TYPE, 1)], + }; + const rowSet: TRowSet = { + startRowOffset: new Int64(0), + rows: [], + columns: [ + { i32Val: { values: [1, 2], nulls: Buffer.from([]) } }, + { stringVal: { values: ['x', 'y'], nulls: Buffer.from([]) } }, + ], + }; + const handler = new JsonResultHandler( + new ClientContextStub(), + new ResultsProviderStub([rowSet], undefined), + { schema, status: { statusCode: TStatusCode.SUCCESS_STATUS } }, + { rowMode: 'array' }, + ); + expect(await handler.fetchNext({ limit: 100 })).to.deep.equal([ + [1, 'x'], + [2, 'y'], + ]); + }); + + it('keeps the default object rows unchanged', async () => { + const handler = new JsonResultHandler( + new ClientContextStub(), + new ResultsProviderStub([duplicateRowSet()], undefined), + { schema: duplicateThriftSchema, status: { statusCode: TStatusCode.SUCCESS_STATUS } }, + ); + expect(await handler.fetchNext({ limit: 100 })).to.deep.equal([{ a: 'E' }, { a: null }]); + }); + + it('matches the column-based compatibility fixture', async () => { + const handler = new JsonResultHandler( + new ClientContextStub(), + new ResultsProviderStub(fixtureColumn.rowSets, undefined), + { schema: fixtureColumn.schema, status: { statusCode: TStatusCode.SUCCESS_STATUS } }, + { rowMode: 'array' }, + ); + const rows = await handler.fetchNext({ limit: 10000 }); + expect(rows).to.deep.equal(toArrays(fixtureColumn.expected, fixtureColumn.schema)); + }); + }); + + describe('ArrowResultConverter', () => { + const duplicateBatch: ArrowBatch = { + batches: [ + arrowIpc(duplicateArrowSchema, [ + [0.23, 0.21], + ['E', null], + ]), + ], + rowCount: 2, + }; + + it('keeps every column of duplicate-named columns in array mode', async () => { + const converter = new ArrowResultConverter( + new ClientContextStub(), + new ResultsProviderStub([duplicateBatch], { batches: [], rowCount: 0 }), + { schema: duplicateThriftSchema }, + { rowMode: 'array' }, + ); + expect(await converter.fetchNext({ limit: 100 })).to.deep.equal([ + [0.23, 'E'], + [0.21, null], + ]); + }); + + it('converts each same-named column with its own type', async () => { + const ts = new Date('2026-01-02T03:04:05.000Z'); + const schema = new Schema([new Field('a', new TimestampMillisecond(), true), new Field('a', new Utf8(), true)]); + const converter = new ArrowResultConverter( + new ClientContextStub(), + new ResultsProviderStub([{ batches: [arrowIpc(schema, [[ts.valueOf()], ['x']])], rowCount: 1 }], { + batches: [], + rowCount: 0, + }), + { schema: { columns: [column('a', TTypeId.TIMESTAMP_TYPE, 1), column('a', TTypeId.STRING_TYPE, 2)] } }, + { rowMode: 'array' }, + ); + expect(await converter.fetchNext({ limit: 100 })).to.deep.equal([[ts, 'x']]); + }); + + it('keeps the default object rows unchanged', async () => { + const converter = new ArrowResultConverter( + new ClientContextStub(), + new ResultsProviderStub([duplicateBatch], { batches: [], rowCount: 0 }), + { schema: duplicateThriftSchema }, + ); + expect(await converter.fetchNext({ limit: 100 })).to.deep.equal([{ a: 'E' }, { a: null }]); + }); + + it('honours the batch row count across multiple record batches', async () => { + const schema = new Schema([new Field('id', new Int32(), true)]); + const thriftSchema: TTableSchema = { columns: [column('id', TTypeId.INT_TYPE, 1)] }; + const converter = new ArrowResultConverter( + new ClientContextStub(), + new ResultsProviderStub( + [ + { batches: [arrowIpc(schema, [[1, 2]])], rowCount: 2 }, + // The server may send more rows than `rowCount` covers; extra rows are dropped. + { batches: [arrowIpc(schema, [[3, 4, 5]])], rowCount: 2 }, + ], + { batches: [], rowCount: 0 }, + ), + { schema: thriftSchema }, + { rowMode: 'array' }, + ); + expect(await converter.fetchNext({ limit: 100 })).to.deep.equal([[1], [2]]); + expect(await converter.fetchNext({ limit: 100 })).to.deep.equal([[3], [4]]); + expect(await converter.hasMore()).to.be.false; + }); + + it('honours preserveBigNumericPrecision', async () => { + const schema = new Schema([ + new Field('n', new ArrowInt64(), true), + new Field('n', new Decimal(2, 20, 128), true), + ]); + const unscaled = new Uint32Array([150, 0, 0, 0]); + const converter = new ArrowResultConverter( + new ClientContextStub(), + new ResultsProviderStub( + [{ batches: [arrowIpc(schema, [[BigInt('9007199254740993')], [unscaled]])], rowCount: 1 }], + { batches: [], rowCount: 0 }, + ), + { schema: { columns: [column('n', TTypeId.BIGINT_TYPE, 1), column('n', TTypeId.DECIMAL_TYPE, 2)] } }, + { rowMode: 'array', preserveBigNumericPrecision: true }, + ); + expect(await converter.fetchNext({ limit: 100 })).to.deep.equal([[BigInt('9007199254740993'), '1.50']]); + }); + + it('still returns null placeholders when row materialization is disabled', async () => { + const converter = new ArrowResultConverter( + new ClientContextStub(), + new ResultsProviderStub([duplicateBatch], { batches: [], rowCount: 0 }), + { schema: duplicateThriftSchema }, + { rowMode: 'array', disableRowMaterialization: true }, + ); + expect(await converter.fetchNext({ limit: 100 })).to.deep.equal([null, null]); + }); + + for (const [name, fixture] of [ + ['without native types', fixtureArrow], + ['with native types', fixtureArrowNT], + ] as const) { + it(`matches the arrow-based compatibility fixture ${name}`, async () => { + const context = new ClientContextStub(); + const converter = new ArrowResultConverter( + context, + new ArrowResultHandler(context, new ResultsProviderStub(fixture.rowSets, undefined), { + arrowSchema: fixture.arrowSchema, + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + }), + { schema: fixture.schema }, + { rowMode: 'array' }, + ); + const rows = (await converter.fetchNext({ limit: 10000 })) as Array>; + const flt = getSchemaColumns(fixture.schema).findIndex((c) => c.columnName === 'flt'); + const fixed = rows.map((row) => + row.map((value, i) => (i === flt ? Number((value as number).toFixed(4)) : value)), + ); + expect(fixed).to.deep.equal(toArrays(fixArrowResult(fixture.expected), fixture.schema)); + }); + } + }); + + describe('duplicate column warning', () => { + it('warns once per result set in object mode', async () => { + const context = new ClientContextStub(); + const spy = sinon.spy(context.logger, 'log'); + const handler = new JsonResultHandler( + context, + new ResultsProviderStub([duplicateRowSet(), duplicateRowSet()], undefined), + { schema: duplicateThriftSchema, status: { statusCode: TStatusCode.SUCCESS_STATUS } }, + ); + await handler.fetchNext({ limit: 100 }); + await handler.fetchNext({ limit: 100 }); + const calls = warnings(context, spy); + expect(calls).to.have.length(1); + expect(calls[0].args[1]).to.contain('"a"'); + expect(calls[0].args[1]).to.contain("rowMode: 'array'"); + + const arrowContext = new ClientContextStub(); + const arrowSpy = sinon.spy(arrowContext.logger, 'log'); + // eslint-disable-next-line no-new + new ArrowResultConverter(arrowContext, new ResultsProviderStub([], { batches: [], rowCount: 0 }), { + schema: duplicateThriftSchema, + }); + expect(warnings(arrowContext, arrowSpy)).to.have.length(1); + }); + + it('does not warn in array mode or for unique column names', async () => { + const context = new ClientContextStub(); + const spy = sinon.spy(context.logger, 'log'); + // eslint-disable-next-line no-new + new JsonResultHandler( + context, + new ResultsProviderStub([], undefined), + { schema: duplicateThriftSchema, status: { statusCode: TStatusCode.SUCCESS_STATUS } }, + { rowMode: 'array' }, + ); + // eslint-disable-next-line no-new + new ArrowResultConverter( + context, + new ResultsProviderStub([], { batches: [], rowCount: 0 }), + { schema: duplicateThriftSchema }, + { rowMode: 'array' }, + ); + // eslint-disable-next-line no-new + new JsonResultHandler(context, new ResultsProviderStub([], undefined), { + schema: fixtureColumn.schema, + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + }); + expect(warnings(context, spy)).to.have.length(0); + }); + }); + + describe('executeStatement', () => { + // RowSetProvider treats a non-empty page as "more rows", so serve one page then an empty one. + function serveRowSet(context: ClientContextStub, results: TRowSet) { + const empty = { + ...results, + columns: results.columns?.map(() => ({ stringVal: { values: [], nulls: Buffer.from([]) } })), + }; + sinon + .stub(context.driver, 'fetchResults') + .resolves({ status: { statusCode: TStatusCode.SUCCESS_STATUS }, hasMoreRows: false, results: empty }) + .onFirstCall() + .resolves({ status: { statusCode: TStatusCode.SUCCESS_STATUS }, hasMoreRows: false, results }); + } + + function createDuplicateColumnSession(context: ClientContextStub) { + context.driver.executeStatementResp.operationHandle!.hasResultSet = true; + context.driver.getResultSetMetadataResp = { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + resultFormat: TSparkRowSetType.COLUMN_BASED_SET, + schema: duplicateThriftSchema, + }; + return createSessionForTest({ + handle: { sessionId: { guid: Buffer.alloc(16), secret: Buffer.alloc(16) } }, + context, + }); + } + + it('returns array rows from the Thrift backend', async () => { + const context = new ClientContextStub(); + const session = createDuplicateColumnSession(context); + serveRowSet(context, duplicateRowSet()); + const operation = await session.executeStatement('SELECT carat AS a, color AS a', { rowMode: 'array' }); + expect(await operation.fetchAll()).to.deep.equal([ + [0.23, 'E'], + [0.21, null], + ]); + }); + + it('keeps object rows by default', async () => { + const context = new ClientContextStub(); + const session = createDuplicateColumnSession(context); + serveRowSet(context, duplicateRowSet()); + const operation = await session.executeStatement('SELECT carat AS a, color AS a'); + expect(await operation.fetchAll()).to.deep.equal([{ a: 'E' }, { a: null }]); + }); + + it('rejects an unsupported rowMode', async () => { + const context = new ClientContextStub(); + const session = createDuplicateColumnSession(context); + const spy = sinon.spy(context.driver, 'executeStatement'); + try { + await session.executeStatement('SELECT 1', { rowMode: 'arrays' as never }); + expect.fail('It should throw an error'); + } catch (error) { + expect(error).to.be.instanceOf(HiveDriverError); + expect((error as Error).message).to.contain('rowMode'); + } + expect(spy.called).to.be.false; + }); + + it('still parses staging responses in array mode', async () => { + const context = new ClientContextStub(); + const session = createDuplicateColumnSession(context); + context.driver.getResultSetMetadataResp = { + status: { statusCode: TStatusCode.SUCCESS_STATUS }, + resultFormat: TSparkRowSetType.COLUMN_BASED_SET, + isStagingOperation: true, + schema: { + columns: [ + column('operation', TTypeId.STRING_TYPE, 1), + column('presignedUrl', TTypeId.STRING_TYPE, 2), + column('headers', TTypeId.STRING_TYPE, 3), + ], + }, + }; + serveRowSet(context, { + startRowOffset: new Int64(0), + rows: [], + columns: [ + { stringVal: { values: ['REMOVE'], nulls: Buffer.from([]) } }, + { stringVal: { values: ['https://example.com/file'], nulls: Buffer.from([]) } }, + { stringVal: { values: ['{}'], nulls: Buffer.from([]) } }, + ], + }); + const remove = sinon.stub(session as any, 'handleStagingRemove').resolves(); + await session.executeStatement('REMOVE "/Volumes/x"', { rowMode: 'array', stagingAllowedLocalPath: '/tmp' }); + expect(remove.calledOnceWith('https://example.com/file', '{}')).to.be.true; + }); + }); + + describe('KernelOperationBackend', () => { + class StatementStub { + public readonly statementId = 'kernel-row-mode'; + + private readonly batches: Buffer[]; + + constructor(private readonly schemaIpc: Buffer, batches: Buffer[]) { + this.batches = [...batches]; + } + + public async fetchNextBatch() { + const ipcBytes = this.batches.shift(); + return ipcBytes ? { ipcBytes } : null; + } + + public schema() { + return { ipcBytes: this.schemaIpc }; + } + + public async cancel() {} + + public async close() {} + } + + function withTypeName(field: Field, typeName: string): Field { + return new Field(field.name, field.type, field.nullable, new Map([['databricks.type_name', typeName]])); + } + + const schema = new Schema([ + withTypeName(new Field('a', new Float64(), true), 'DOUBLE'), + withTypeName(new Field('a', new Utf8(), true), 'STRING'), + ]); + + it('returns array rows when rowMode is array', async () => { + const backend = new KernelOperationBackend({ + statement: new StatementStub(arrowIpc(schema, [[], []]), [ + arrowIpc(schema, [ + [0.23, 0.21], + ['E', null], + ]), + ]), + context: new ClientContextStub(), + rowMode: 'array', + }); + expect(await backend.fetchChunk({ limit: 100 })).to.deep.equal([ + [0.23, 'E'], + [0.21, null], + ]); + }); + + it('keeps object rows by default', async () => { + const backend = new KernelOperationBackend({ + statement: new StatementStub(arrowIpc(schema, [[], []]), [ + arrowIpc(schema, [ + [0.23, 0.21], + ['E', null], + ]), + ]), + context: new ClientContextStub(), + }); + expect(await backend.fetchChunk({ limit: 100 })).to.deep.equal([{ a: 'E' }, { a: null }]); + }); + }); +}); From 97412c98f3c5dbdef55aab52dc298e7269832e15 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 23:25:05 +0000 Subject: [PATCH 2/3] test: add e2e coverage for rowMode Cover duplicate-alias rows in array mode on Thrift column-based, inline Arrow (LZ4) and CloudFetch results, staging statements in array mode, and the kernel and Thrift backends in the kernel parity spec. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- tests/e2e/kernel/results-e2e.test.ts | 23 +++++- tests/e2e/row_mode.test.ts | 111 +++++++++++++++++++++++++++ 2 files changed, 132 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/row_mode.test.ts diff --git a/tests/e2e/kernel/results-e2e.test.ts b/tests/e2e/kernel/results-e2e.test.ts index 5f5969d7..1787972c 100644 --- a/tests/e2e/kernel/results-e2e.test.ts +++ b/tests/e2e/kernel/results-e2e.test.ts @@ -18,6 +18,7 @@ import { expect } from 'chai'; import { DBSQLClient } from '../../../lib'; import { ConnectionOptions } from '../../../lib/contracts/IDBSQLClient'; import { InternalConnectionOptions } from '../../../lib/contracts/InternalConnectionOptions'; +import { ExecuteStatementOptions } from '../../../lib/contracts/IDBSQLSession'; // Integration suite: connect through both backends, run a probe query, // and assert byte-identical row output (the M0 parity gate). Requires @@ -44,7 +45,12 @@ function readSecrets(): PecoSecrets | null { return { host, path, token }; } -async function fetchProbeRows(useKernel: boolean, secrets: PecoSecrets): Promise>> { +async function fetchProbeRows( + useKernel: boolean, + secrets: PecoSecrets, + query: string = PROBE_QUERY, + options: ExecuteStatementOptions = {}, +): Promise>> { const client = new DBSQLClient(); await client.connect({ host: secrets.host, @@ -55,7 +61,7 @@ async function fetchProbeRows(useKernel: boolean, secrets: PecoSecrets): Promise try { const session = await client.openSession(); try { - const operation = await session.executeStatement(PROBE_QUERY); + const operation = await session.executeStatement(query, options); try { const rows = (await operation.fetchAll()) as Array>; return rows; @@ -125,4 +131,17 @@ describe('kernel results end-to-end (pecotesting parity gate)', function suite() const thriftRows = await fetchProbeRows(false, secrets as PecoSecrets); expect(kernelRows.map(canonical)).to.deep.equal(thriftRows.map(canonical)); }); + + it('keeps same-named columns with rowMode array on both backends (#118)', async () => { + const query = + "SELECT carat AS a, color AS a FROM VALUES (0.23, 'E'), (0.21, 'I') AS t(carat, color) ORDER BY carat DESC"; + for (const useKernel of [true, false]) { + // eslint-disable-next-line no-await-in-loop + const rows = await fetchProbeRows(useKernel, secrets as PecoSecrets, query, { rowMode: 'array' }); + expect(rows, useKernel ? 'kernel' : 'thrift').to.deep.equal([ + [0.23, 'E'], + [0.21, 'I'], + ]); + } + }); }); diff --git a/tests/e2e/row_mode.test.ts b/tests/e2e/row_mode.test.ts new file mode 100644 index 00000000..665449ea --- /dev/null +++ b/tests/e2e/row_mode.test.ts @@ -0,0 +1,111 @@ +import { expect } from 'chai'; +import sinon from 'sinon'; +import fs from 'fs'; +import os from 'os'; +import path from 'path'; +import * as uuid from 'uuid'; +import { DBSQLClient } from '../../lib'; +import { ClientConfig } from '../../lib/contracts/IClientContext'; +import IDBSQLSession from '../../lib/contracts/IDBSQLSession'; +import ArrowResultHandler from '../../lib/result/ArrowResultHandler'; +import CloudFetchResultHandler from '../../lib/result/CloudFetchResultHandler'; +import JsonResultHandler from '../../lib/result/JsonResultHandler'; + +import config from './utils/config'; + +// From #118: same-named columns overwrite each other in object rows. +const DUPLICATE_ALIAS_QUERY = + "SELECT carat AS a, color AS a FROM VALUES (0.23, 'E'), (0.21, 'I') AS t(carat, color) ORDER BY carat DESC"; + +async function openSession(customConfig: Partial = {}): Promise { + const client = new DBSQLClient(); + const clientConfig = client.getConfig(); + sinon.stub(client, 'getConfig').returns({ ...clientConfig, ...customConfig }); + const connection = await client.connect({ host: config.host, path: config.path, token: config.token }); + return connection.openSession({ initialCatalog: config.catalog, initialSchema: config.schema }); +} + +// Whether the operation's result pipeline (converter or its source) is an instance of `handlerClass`. +async function usesResultHandler(operation: unknown, handlerClass: new (...args: any[]) => unknown): Promise { + // @ts-expect-error TS2339: getResultHandler is not part of IOperation + const resultHandler = await operation.getResultHandler(); + return [resultHandler.source, resultHandler.source.source].some((source) => source instanceof handlerClass); +} + +describe('rowMode', () => { + for (const [name, clientConfig, handlerClass] of [ + ['column-based results', { arrowEnabled: false }, JsonResultHandler], + ['inline Arrow results', { arrowEnabled: true, useLZ4Compression: true }, ArrowResultHandler], + ] as const) { + it(`keeps every same-named column with ${name}`, async () => { + const session = await openSession(clientConfig); + + const arrayOperation = await session.executeStatement(DUPLICATE_ALIAS_QUERY, { + useCloudFetch: false, + rowMode: 'array', + }); + expect(await usesResultHandler(arrayOperation, handlerClass)).to.be.true; + expect(await arrayOperation.fetchAll()).to.deep.equal([ + [0.23, 'E'], + [0.21, 'I'], + ]); + await arrayOperation.close(); + + const objectOperation = await session.executeStatement(DUPLICATE_ALIAS_QUERY, { useCloudFetch: false }); + expect(await objectOperation.fetchAll()).to.deep.equal([{ a: 'E' }, { a: 'I' }]); + await objectOperation.close(); + + await session.close(); + }); + } + + it('keeps every same-named column with CloudFetch results', async () => { + const session = await openSession({ cloudFetchConcurrentDownloads: 5 }); + const rowCount = 2000000; // result has to be quite big to enable CloudFetch + const operation = await session.executeStatement( + `SELECT id AS a, CAST(id AS STRING) AS a FROM range(0, ${rowCount})`, + { maxRows: null, useCloudFetch: true, rowMode: 'array' }, + ); + await operation.finished(); + expect(await usesResultHandler(operation, CloudFetchResultHandler)).to.be.true; + + let fetched = 0; + do { + // eslint-disable-next-line no-await-in-loop + const chunk = (await operation.fetchChunk({ maxRows: 100000, disableBuffering: true })) as Array< + [number, string] + >; + for (const [id, text] of chunk) { + expect(text).to.equal(String(id)); + } + fetched += chunk.length; + } while (await operation.hasMoreRows()); // eslint-disable-line no-await-in-loop + expect(fetched).to.equal(rowCount); + + await operation.close(); + await session.close(); + }); + + it('runs staging statements in array mode', async () => { + const { catalog, schema, volume } = config; + const localPath = fs.mkdtempSync(path.join(os.tmpdir(), 'databricks-sql-tests-')); + try { + const session = await openSession(); + const stagingFileName = `/Volumes/${catalog}/${schema}/${volume}/${uuid.v4()}.csv`; + const localFile = path.join(localPath, `${uuid.v4()}.csv`); + const options = { stagingAllowedLocalPath: [localPath], rowMode: 'array' as const }; + + fs.writeFileSync(localFile, 'Hello World!'); + await session.executeStatement(`PUT '${localFile}' INTO '${stagingFileName}' OVERWRITE`, options); + fs.rmSync(localFile); + + await session.executeStatement(`GET '${stagingFileName}' TO '${localFile}'`, options); + expect(fs.readFileSync(localFile).toString()).to.equal('Hello World!'); + + await session.executeStatement(`REMOVE '${stagingFileName}'`, options); + await session.close(); + } finally { + fs.rmSync(localPath, { recursive: true, force: true }); + } + }); +}); From 6334f9c732a777bfb268c5e56133f5fa72ec47dc Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Wed, 11 Nov 2026 23:36:39 +0000 Subject: [PATCH 3/3] feat: export RowMode from the package entry point Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- lib/index.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/index.ts b/lib/index.ts index a2b90c9c..2cffc3ea 100644 --- a/lib/index.ts +++ b/lib/index.ts @@ -22,6 +22,7 @@ import { LogLevel } from './contracts/IDBSQLLogger'; // Re-export types for TypeScript users export type { default as ITokenProvider } from './connection/auth/tokenProvider/ITokenProvider'; +export type { RowMode } from './contracts/IDBSQLSession'; export { LOG_FILE_UNAVAILABLE_WARNING_CODE } from './DBSQLLogger'; export { CircuitBreakerOpenError, CIRCUIT_BREAKER_OPEN_CODE } from './telemetry/CircuitBreaker';