Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/appkit-ui/src/react/charts/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import type { Table } from "apache-arrow";
// ============================================================================

/** Supported data formats for analytics queries */
export type DataFormat = "json" | "arrow" | "auto";
export type DataFormat = "json" | "arrow" | "arrow_stream" | "auto";

/** Chart orientation */
export type Orientation = "vertical" | "horizontal";
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ describe("useChartData", () => {
);
});

test("auto-selects JSON by default when no heuristics match", () => {
test("auto-selects ARROW_STREAM by default when no heuristics match", () => {
mockUseAnalyticsQuery.mockReturnValue({
data: [],
loading: false,
Expand All @@ -223,11 +223,11 @@ describe("useChartData", () => {
expect(mockUseAnalyticsQuery).toHaveBeenCalledWith(
"test",
{ limit: 100 },
expect.objectContaining({ format: "JSON" }),
expect.objectContaining({ format: "ARROW_STREAM" }),
);
});

test("defaults to auto format (JSON) when format is not specified", () => {
test("defaults to auto format (ARROW_STREAM) when format is not specified", () => {
mockUseAnalyticsQuery.mockReturnValue({
data: [],
loading: false,
Expand All @@ -243,7 +243,7 @@ describe("useChartData", () => {
expect(mockUseAnalyticsQuery).toHaveBeenCalledWith(
"test",
undefined,
expect.objectContaining({ format: "JSON" }),
expect.objectContaining({ format: "ARROW_STREAM" }),
);
});
});
Expand Down
8 changes: 5 additions & 3 deletions packages/appkit-ui/src/react/hooks/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import type { Table } from "apache-arrow";
// ============================================================================

/** Supported response formats for analytics queries */
export type AnalyticsFormat = "JSON" | "ARROW";
export type AnalyticsFormat = "JSON" | "ARROW" | "ARROW_STREAM";

/**
* Typed Arrow Table - preserves row type information for type inference.
Expand All @@ -32,8 +32,10 @@ export interface TypedArrowTable<
// ============================================================================

/** Options for configuring an analytics SSE query */
export interface UseAnalyticsQueryOptions<F extends AnalyticsFormat = "JSON"> {
/** Response format - "JSON" returns typed arrays, "ARROW" returns TypedArrowTable */
export interface UseAnalyticsQueryOptions<
F extends AnalyticsFormat = "ARROW_STREAM",
> {
/** Response format - "ARROW_STREAM" (default) uses inline Arrow, "JSON" returns typed arrays, "ARROW" uses external links */
format?: F;

/** Maximum size of serialized parameters in bytes */
Expand Down
4 changes: 2 additions & 2 deletions packages/appkit-ui/src/react/hooks/use-analytics-query.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,13 @@ function getArrowStreamUrl(id: string) {
export function useAnalyticsQuery<
T = unknown,
K extends QueryKey = QueryKey,
F extends AnalyticsFormat = "JSON",
F extends AnalyticsFormat = "ARROW_STREAM",
>(
queryKey: K,
parameters?: InferParams<K> | null,
options: UseAnalyticsQueryOptions<F> = {} as UseAnalyticsQueryOptions<F>,
): UseAnalyticsQueryResult<InferResultByFormat<T, K, F>> {
const format = options?.format ?? "JSON";
const format = options?.format ?? "ARROW_STREAM";
const maxParametersSize = options?.maxParametersSize ?? 100 * 1024;
const autoStart = options?.autoStart ?? true;

Expand Down
7 changes: 4 additions & 3 deletions packages/appkit-ui/src/react/hooks/use-chart-data.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,11 @@ export interface UseChartDataResult {
function resolveFormat(
format: DataFormat,
parameters?: Record<string, unknown>,
): "JSON" | "ARROW" {
): "JSON" | "ARROW" | "ARROW_STREAM" {
// Explicit format selection
if (format === "json") return "JSON";
if (format === "arrow") return "ARROW";
if (format === "arrow_stream") return "ARROW_STREAM";

// Auto-selection heuristics
if (format === "auto") {
Expand All @@ -72,10 +73,10 @@ function resolveFormat(
return "ARROW";
}

return "JSON";
return "ARROW_STREAM";
}

return "JSON";
return "ARROW_STREAM";
}

// ============================================================================
Expand Down
1 change: 1 addition & 0 deletions packages/appkit/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@
"@opentelemetry/sdk-trace-base": "2.6.0",
"@opentelemetry/semantic-conventions": "1.38.0",
"@types/semver": "7.7.1",
"apache-arrow": "21.1.0",
"dotenv": "16.6.1",
"express": "4.22.0",
"obug": "2.1.1",
Expand Down
38 changes: 37 additions & 1 deletion packages/appkit/src/connectors/sql-warehouse/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import {
type sql,
type WorkspaceClient,
} from "@databricks/sdk-experimental";
import { tableFromIPC } from "apache-arrow";
import type { TelemetryOptions } from "shared";
import {
AppKitError,
Expand Down Expand Up @@ -393,7 +394,20 @@ export class SQLWarehouseConnector {

private _transformDataArray(response: sql.StatementResponse) {
if (response.manifest?.format === "ARROW_STREAM") {
return this.updateWithArrowStatus(response);
const result = response.result as any;

// Inline Arrow: some warehouses return base64 Arrow IPC in `attachment`.
if (result?.attachment) {
return this._transformArrowAttachment(response, result.attachment);
}

// Inline data_array: fall through to the row transform below.
if (result?.data_array) {
// Fall through.
} else {
// External links: data fetched separately via statement_id.
return this.updateWithArrowStatus(response);
}
}

if (!response.result?.data_array || !response.manifest?.schema?.columns) {
Expand Down Expand Up @@ -439,6 +453,28 @@ export class SQLWarehouseConnector {
};
}

/**
* Decode a base64 Arrow IPC attachment into row objects.
* Some serverless warehouses return inline results as Arrow IPC in
* `result.attachment` rather than `result.data_array`.
*/
private _transformArrowAttachment(
response: sql.StatementResponse,
attachment: string,
) {
const buf = Buffer.from(attachment, "base64");
const table = tableFromIPC(buf);
const data = table.toArray().map((row) => row.toJSON());
const { attachment: _att, ...restResult } = response.result as any;
return {
...response,
result: {
...restResult,
data,
},
};
}

private updateWithArrowStatus(response: sql.StatementResponse): {
result: { statement_id: string; status: sql.StatementStatus };
} {
Expand Down
Loading
Loading