Skip to content

Commit 6e34164

Browse files
fix: handle ARROW_STREAM attachment in type generator
When serverless warehouses return ARROW_STREAM format, the DESCRIBE QUERY result comes as an inline base64 Arrow IPC attachment rather than data_array. This caused convertToQueryType to generate empty types {}. Add a fallback that decodes the Arrow IPC attachment schema to extract column names and types when data_array is empty. Co-authored-by: Isaac Signed-off-by: James Broadhead <jamesbroadhead@gmail.com>
1 parent 20dca67 commit 6e34164

2 files changed

Lines changed: 73 additions & 2 deletions

File tree

packages/appkit/src/type-generator/query-registry.ts

Lines changed: 71 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import fs from "node:fs/promises";
22
import path from "node:path";
33
import { WorkspaceClient } from "@databricks/sdk-experimental";
4+
import { tableFromIPC } from "apache-arrow";
45
import pc from "picocolors";
56
import { createLogger } from "../logging/logger";
67
import { CACHE_VERSION, hashSQL, loadCache, saveCache } from "./cache";
@@ -78,18 +79,85 @@ function formatParametersType(sql: string): string {
7879
: "Record<string, never>";
7980
}
8081

82+
/**
83+
* Map Arrow DataType IDs to Databricks SQL type names.
84+
* Arrow type IDs come from the Arrow spec (apache-arrow TypeId enum).
85+
* We only need to cover the types that DESCRIBE QUERY can return.
86+
*/
87+
function arrowTypeToSqlName(arrowType: { typeId: number }): string {
88+
switch (arrowType.typeId) {
89+
case 1: // Bool
90+
return "BOOLEAN";
91+
case 2: // Int (covers TINYINT, SMALLINT, INT, BIGINT depending on bitWidth)
92+
return "INT";
93+
case 3: // Float (covers FLOAT, DOUBLE)
94+
return "DOUBLE";
95+
case 4: // Decimal
96+
return "DECIMAL";
97+
case 5: // Utf8
98+
return "STRING";
99+
case 6: // Binary
100+
return "BINARY";
101+
case 7: // FixedSizeBinary
102+
return "BINARY";
103+
case 8: // Date
104+
return "DATE";
105+
case 10: // Timestamp
106+
return "TIMESTAMP";
107+
case 12: // List
108+
return "ARRAY";
109+
case 14: // Struct
110+
return "STRUCT";
111+
case 15: // Map
112+
return "MAP";
113+
default:
114+
return "STRING";
115+
}
116+
}
117+
118+
/**
119+
* Decode a base64 Arrow IPC attachment and extract column metadata.
120+
* Returns the same shape as rows parsed from DESCRIBE QUERY data_array.
121+
*/
122+
function columnsFromArrowAttachment(
123+
attachment: string,
124+
): Array<{ name: string; type_name: string; comment: string | undefined }> {
125+
const buf = Buffer.from(attachment, "base64");
126+
const table = tableFromIPC(buf);
127+
return table.schema.fields.map((field) => ({
128+
name: field.name,
129+
type_name: arrowTypeToSqlName(field.type),
130+
comment: undefined,
131+
}));
132+
}
133+
81134
export function convertToQueryType(
82135
result: DatabricksStatementExecutionResponse,
83136
sql: string,
84137
queryName: string,
85138
): { type: string; hasResults: boolean } {
86139
const dataRows = result.result?.data_array || [];
87-
const columns = dataRows.map((row) => ({
140+
let columns = dataRows.map((row) => ({
88141
name: row[0] || "",
89142
type_name: row[1]?.toUpperCase() || "STRING",
90143
comment: row[2] || undefined,
91144
}));
92145

146+
// Fallback: serverless warehouses may return ARROW_STREAM format with an
147+
// inline base64 attachment instead of data_array. Decode the Arrow IPC
148+
// schema to extract column names and types.
149+
if (columns.length === 0 && result.result?.attachment) {
150+
logger.debug("data_array empty, decoding Arrow IPC attachment for schema");
151+
try {
152+
columns = columnsFromArrowAttachment(result.result.attachment);
153+
} catch (err) {
154+
logger.warn(
155+
"Failed to decode Arrow IPC attachment: %s",
156+
err instanceof Error ? err.message : String(err),
157+
);
158+
}
159+
}
160+
93161
const paramsType = formatParametersType(sql);
94162

95163
// generate result fields with JSDoc
@@ -277,10 +345,11 @@ export async function generateQueriesFromDescribe(
277345
);
278346

279347
logger.debug(
280-
"DESCRIBE result for %s: state=%s, rows=%d",
348+
"DESCRIBE result for %s: state=%s, rows=%d, hasAttachment=%s",
281349
queryName,
282350
result.status.state,
283351
result.result?.data_array?.length ?? 0,
352+
!!result.result?.attachment,
284353
);
285354

286355
if (result.status.state === "FAILED") {

packages/appkit/src/type-generator/types.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ export interface DatabricksStatementExecutionResponse {
1212
};
1313
result?: {
1414
data_array?: (string | null)[][];
15+
/** Base64-encoded Arrow IPC bytes (returned by serverless warehouses using ARROW_STREAM format) */
16+
attachment?: string;
1517
};
1618
}
1719

0 commit comments

Comments
 (0)