Add Arrow field mapping header for native load exports (#24)

Author: acmeguy

services/cubejs/src/routes/loadExport.js

@@ -40,6 +40,10 @@ const ARROW_HEADERS = {
   "Content-Type": "application/vnd.apache.arrow.stream",
   "Content-Disposition": 'attachment; filename="query-result.arrow"',
 };
+const ARROW_FIELD_MAPPING_HEADER = "X-Synmetrix-Arrow-Field-Mapping";
+const ARROW_FIELD_MAPPING_ENCODING_HEADER =
+  "X-Synmetrix-Arrow-Field-Mapping-Encoding";
+const ARROW_FIELD_MAPPING_ENCODING = "base64url-json";
 const CSV_STREAM_CHUNK_SIZE = 128 * 1024;
 const CLICKHOUSE_NULL_TOKEN_RE = /(?<=,|^)\\N(?=,|\r?\n|$)/g;
 
@@ -167,15 +171,71 @@ function getAliasNameToMember(plan) {
   return plan?.streamingQuery?.aliasNameToMember || null;
 }
 
-function canUseNativeArrowPassthrough(plan) {
+function getChangedAliasNameToMember(plan) {
   const aliasNameToMember = getAliasNameToMember(plan);
-  if (!aliasNameToMember || Object.keys(aliasNameToMember).length === 0) {
-    return true;
+  if (!aliasNameToMember) return null;
+
+  const changedAliasNameToMember = Object.fromEntries(
+    Object.entries(aliasNameToMember).filter(([alias, member]) => alias !== member)
+  );
+
+  return Object.keys(changedAliasNameToMember).length > 0
+    ? changedAliasNameToMember
+    : null;
+}
+
+function getResponseHeader(res, name) {
+  if (typeof res.getHeader === "function") {
+    const value = res.getHeader(name);
+    if (value != null) return value;
+  }
+
+  if (typeof res.get === "function") {
+    const value = res.get(name);
+    if (value != null) return value;
   }
 
-  return Object.entries(aliasNameToMember).every(
-    ([alias, member]) => alias === member
+  return res.headers?.[name] ?? res.headers?.[name.toLowerCase()];
+}
+
+function appendResponseHeaderValues(res, name, values) {
+  const existing = getResponseHeader(res, name);
+  const currentValues = Array.isArray(existing)
+    ? existing.flatMap((value) => String(value).split(","))
+    : typeof existing === "string"
+      ? existing.split(",")
+      : [];
+  const normalized = new Set(
+    currentValues.map((value) => value.trim()).filter(Boolean).map((value) => value.toLowerCase())
   );
+  const mergedValues = currentValues.map((value) => value.trim()).filter(Boolean);
+
+  for (const value of values) {
+    const normalizedValue = value.toLowerCase();
+    if (!normalized.has(normalizedValue)) {
+      normalized.add(normalizedValue);
+      mergedValues.push(value);
+    }
+  }
+
+  res.set(name, mergedValues.join(", "));
+}
+
+function setNativeArrowFieldMappingHeaders(res, plan) {
+  const aliasNameToMember = getChangedAliasNameToMember(plan);
+  if (!aliasNameToMember) return;
+
+  const encodedMapping = Buffer.from(
+    JSON.stringify(aliasNameToMember),
+    "utf8"
+  ).toString("base64url");
+
+  res.set(ARROW_FIELD_MAPPING_HEADER, encodedMapping);
+  res.set(ARROW_FIELD_MAPPING_ENCODING_HEADER, ARROW_FIELD_MAPPING_ENCODING);
+  appendResponseHeaderValues(res, "Access-Control-Expose-Headers", [
+    ARROW_FIELD_MAPPING_HEADER,
+    ARROW_FIELD_MAPPING_ENCODING_HEADER,
+  ]);
 }
 
 function rewriteNativeCsvHeader(line, aliasNameToMember) {
@@ -416,7 +476,6 @@ async function tryHandleLoadExport(req, res, cubejs, query, format) {
     if (
       format === "arrow"
       && plan.capabilities.nativeArrowPassthrough
-      && canUseNativeArrowPassthrough(plan)
       && isClickHouseContext(plan.context.securityContext)
     ) {
       const nativeQuery = await prepareNativeClickHouseExport(plan);
@@ -425,6 +484,7 @@ async function tryHandleLoadExport(req, res, cubejs, query, format) {
           securityContext: plan.context.securityContext,
         });
         res.set(ARROW_HEADERS);
+        setNativeArrowFieldMappingHeaders(res, plan);
         await executeNativeClickHouseArrow(
           res,
           nativeQuery.query,

services/cubejs/test/loadExport.test.js

@@ -5,6 +5,10 @@ import { tableFromIPC } from "apache-arrow";
 
 import { maybeHandleLoadExport } from "../src/routes/loadExport.js";
 
+function decodeArrowFieldMappingHeader(value) {
+  return JSON.parse(Buffer.from(value, "base64url").toString("utf8"));
+}
+
 describe("maybeHandleLoadExport", () => {
   it("streams CSV through the semantic export path when native passthrough is unavailable", async () => {
     const req = createRequest({
@@ -99,7 +103,7 @@ describe("maybeHandleLoadExport", () => {
     );
   });
 
-  it("streams Arrow through the semantic export path when native passthrough is unsafe", async () => {
+  it("streams Arrow through the native ClickHouse path even when aliases are not semantic member names", async () => {
     const req = createRequest({
       format: "arrow",
       query: {
@@ -145,7 +149,7 @@ describe("maybeHandleLoadExport", () => {
         values: [],
       },
       driver: createClickHouseDriver({
-        arrowChunks: [Buffer.from("native-arrow-should-not-be-used")],
+        arrowChunks: [Buffer.from("native-arrow-stream-binary")],
       }),
     });
 
@@ -157,15 +161,25 @@ describe("maybeHandleLoadExport", () => {
       res.headers["Content-Type"],
       "application/vnd.apache.arrow.stream"
     );
-
-    const parsed = tableFromIPC(res.binaryOutput);
-    assert.deepEqual(parsed.toArray().map((row) => row.toJSON()), [
+    assert.deepEqual(
+      decodeArrowFieldMappingHeader(
+        res.headers["X-Synmetrix-Arrow-Field-Mapping"]
+      ),
       {
-        "Orders.city": "Reykjavik",
-        "Orders.count": 2,
-        "Orders.createdAt.day": Date.parse("2024-01-01T00:00:00.000Z"),
-      },
-    ]);
+        city_alias: "Orders.city",
+        count_alias: "Orders.count",
+        created_day_alias: "Orders.createdAt.day",
+      }
+    );
+    assert.equal(
+      res.headers["X-Synmetrix-Arrow-Field-Mapping-Encoding"],
+      "base64url-json"
+    );
+    assert.match(
+      res.headers["Access-Control-Expose-Headers"],
+      /X-Synmetrix-Arrow-Field-Mapping/
+    );
+    assert.deepEqual(res.binaryOutput, Buffer.from("native-arrow-stream-binary"));
   });
 
   it("streams Arrow through the native ClickHouse path when aliases are already semantic", async () => {
@@ -208,9 +222,87 @@ describe("maybeHandleLoadExport", () => {
       res.headers["Content-Type"],
       "application/vnd.apache.arrow.stream"
     );
+    assert.equal(res.headers["X-Synmetrix-Arrow-Field-Mapping"], undefined);
     assert.deepEqual(res.binaryOutput, nativeArrowBuffer);
   });
 
+  it("streams Arrow through the semantic export path when native ClickHouse export is unavailable", async () => {
+    const req = createRequest({
+      format: "arrow",
+      query: {
+        dimensions: ["Orders.city"],
+        measures: ["Orders.count"],
+        timeDimensions: [{ dimension: "Orders.createdAt", granularity: "day" }],
+      },
+    });
+    const res = new MockResponse();
+
+    const cubejs = createMockCube({
+      dbType: "clickhouse",
+      normalizedQuery: req.body.query,
+      sqlQuery: {
+        sql: ["SELECT city_alias, count_alias, created_day_alias FROM orders", []],
+        aliasNameToMember: {
+          city_alias: "Orders.city",
+          count_alias: "Orders.count",
+          created_day_alias: "Orders.createdAt.day",
+        },
+      },
+      metaConfig: createMetaConfig({
+        measures: [{ name: "Orders.count", type: "count" }],
+        dimensions: [
+          { name: "Orders.city", type: "string" },
+          {
+            name: "Orders.createdAt",
+            type: "time",
+            granularities: [{ name: "day", title: "day", interval: "1 day" }],
+          },
+          { name: "Orders.createdAt.day", type: "time" },
+        ],
+      }),
+      streamRows: [
+        {
+          "Orders.city": "Reykjavik",
+          "Orders.count": "2",
+          "Orders.createdAt.day": "2024-01-01T00:00:00.000Z",
+        },
+      ],
+      nativePreAggs: {
+        preAggregationsTablesToTempTables: [
+          [
+            "orders_rollup",
+            {
+              lambdaTable: { name: "orders_lambda" },
+            },
+          ],
+        ],
+        values: [],
+      },
+      driver: createClickHouseDriver({
+        arrowChunks: [Buffer.from("native-arrow-should-not-be-used")],
+      }),
+    });
+
+    await maybeHandleLoadExport(req, res, () => {
+      throw new Error("next should not be called");
+    }, cubejs);
+
+    assert.equal(
+      res.headers["Content-Type"],
+      "application/vnd.apache.arrow.stream"
+    );
+    assert.equal(res.headers["X-Synmetrix-Arrow-Field-Mapping"], undefined);
+
+    const parsed = tableFromIPC(res.binaryOutput);
+    assert.deepEqual(parsed.toArray().map((row) => row.toJSON()), [
+      {
+        "Orders.city": "Reykjavik",
+        "Orders.count": 2,
+        "Orders.createdAt.day": Date.parse("2024-01-01T00:00:00.000Z"),
+      },
+    ]);
+  });
+
   it("falls through to the buffered path when semantic streaming is unavailable", async () => {
     const req = createRequest({
       format: "csv",