feat: add output_schema method to queries (#2717)

This is a helper utility I need for some of my data loader work. It
makes it easy to see the output schema even when a `select` has been
applied.

nodejs/__test__/query.test.ts

@@ -0,0 +1,111 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+import * as tmp from "tmp";
+
+import { type Table, connect } from "../lancedb";
+import {
+  Field,
+  FixedSizeList,
+  Float32,
+  Int64,
+  Schema,
+  Utf8,
+  makeArrowTable,
+} from "../lancedb/arrow";
+import { Index } from "../lancedb/indices";
+
+describe("Query outputSchema", () => {
+  let tmpDir: tmp.DirResult;
+  let table: Table;
+
+  beforeEach(async () => {
+    tmpDir = tmp.dirSync({ unsafeCleanup: true });
+    const db = await connect(tmpDir.name);
+
+    // Create table with explicit schema to ensure proper types
+    const schema = new Schema([
+      new Field("a", new Int64(), true),
+      new Field("text", new Utf8(), true),
+      new Field(
+        "vec",
+        new FixedSizeList(2, new Field("item", new Float32())),
+        true,
+      ),
+    ]);
+
+    const data = makeArrowTable(
+      [
+        { a: 1n, text: "foo", vec: [1, 2] },
+        { a: 2n, text: "bar", vec: [3, 4] },
+        { a: 3n, text: "baz", vec: [5, 6] },
+      ],
+      { schema },
+    );
+    table = await db.createTable("test", data);
+  });
+
+  afterEach(() => {
+    tmpDir.removeCallback();
+  });
+
+  it("should return schema for plain query", async () => {
+    const schema = await table.query().outputSchema();
+
+    expect(schema.fields.length).toBe(3);
+    expect(schema.fields.map((f) => f.name)).toEqual(["a", "text", "vec"]);
+    expect(schema.fields[0].type.toString()).toBe("Int64");
+    expect(schema.fields[1].type.toString()).toBe("Utf8");
+  });
+
+  it("should return schema with dynamic projection", async () => {
+    const schema = await table.query().select({ bl: "a * 2" }).outputSchema();
+
+    expect(schema.fields.length).toBe(1);
+    expect(schema.fields[0].name).toBe("bl");
+    expect(schema.fields[0].type.toString()).toBe("Int64");
+  });
+
+  it("should return schema for vector search with _distance column", async () => {
+    const schema = await table
+      .vectorSearch([1, 2])
+      .select(["a"])
+      .outputSchema();
+
+    expect(schema.fields.length).toBe(2);
+    expect(schema.fields.map((f) => f.name)).toEqual(["a", "_distance"]);
+    expect(schema.fields[0].type.toString()).toBe("Int64");
+    expect(schema.fields[1].type.toString()).toBe("Float32");
+  });
+
+  it("should return schema for FTS search", async () => {
+    await table.createIndex("text", { config: Index.fts() });
+
+    const schema = await table
+      .search("foo", "fts")
+      .select(["a"])
+      .outputSchema();
+
+    // FTS search includes _score column in addition to selected columns
+    expect(schema.fields.length).toBe(2);
+    expect(schema.fields.map((f) => f.name)).toContain("a");
+    expect(schema.fields.map((f) => f.name)).toContain("_score");
+    const aField = schema.fields.find((f) => f.name === "a");
+    expect(aField?.type.toString()).toBe("Int64");
+  });
+
+  it("should return schema for take query", async () => {
+    const schema = await table.takeOffsets([0]).select(["text"]).outputSchema();
+
+    expect(schema.fields.length).toBe(1);
+    expect(schema.fields[0].name).toBe("text");
+    expect(schema.fields[0].type.toString()).toBe("Utf8");
+  });
+
+  it("should return full schema when no select is specified", async () => {
+    const schema = await table.query().outputSchema();
+
+    // Should return all columns
+    expect(schema.fields.length).toBe(3);
+  });
+});

nodejs/lancedb/query.ts

@@ -326,6 +326,25 @@ export class QueryBase<
       return this.inner.analyzePlan();
     }
   }
+
+  /**
+   * Returns the schema of the output that will be returned by this query.
+   *
+   * This can be used to inspect the types and names of the columns that will be
+   * returned by the query before executing it.
+   *
+   * @returns An Arrow Schema describing the output columns.
+   */
+  async outputSchema(): Promise<import("./arrow").Schema> {
+    let schemaBuffer: Buffer;
+    if (this.inner instanceof Promise) {
+      schemaBuffer = await this.inner.then((inner) => inner.outputSchema());
+    } else {
+      schemaBuffer = await this.inner.outputSchema();
+    }
+    const schema = tableFromIPC(schemaBuffer).schema;
+    return schema;
+  }
 }
 
 export class StandardQueryBase<

nodejs/src/query.rs

@@ -22,7 +22,7 @@ use crate::error::NapiErrorExt;
 use crate::iterator::RecordBatchIterator;
 use crate::rerankers::Reranker;
 use crate::rerankers::RerankerCallbacks;
-use crate::util::parse_distance_type;
+use crate::util::{parse_distance_type, schema_to_buffer};
 
 #[napi]
 pub struct Query {
@@ -88,6 +88,12 @@ impl Query {
         self.inner = self.inner.clone().with_row_id();
     }
 
+    #[napi(catch_unwind)]
+    pub async fn output_schema(&self) -> napi::Result<Buffer> {
+        let schema = self.inner.output_schema().await.default_error()?;
+        schema_to_buffer(&schema)
+    }
+
     #[napi(catch_unwind)]
     pub async fn execute(
         &self,
@@ -273,6 +279,12 @@ impl VectorQuery {
             .rerank(Arc::new(Reranker::new(callbacks)));
     }
 
+    #[napi(catch_unwind)]
+    pub async fn output_schema(&self) -> napi::Result<Buffer> {
+        let schema = self.inner.output_schema().await.default_error()?;
+        schema_to_buffer(&schema)
+    }
+
     #[napi(catch_unwind)]
     pub async fn execute(
         &self,
@@ -346,6 +358,12 @@ impl TakeQuery {
         self.inner = self.inner.clone().with_row_id();
     }
 
+    #[napi(catch_unwind)]
+    pub async fn output_schema(&self) -> napi::Result<Buffer> {
+        let schema = self.inner.output_schema().await.default_error()?;
+        schema_to_buffer(&schema)
+    }
+
     #[napi(catch_unwind)]
     pub async fn execute(
         &self,

nodejs/src/table.rs

@@ -3,7 +3,6 @@
 
 use std::collections::HashMap;
 
-use arrow_ipc::writer::FileWriter;
 use lancedb::ipc::ipc_file_to_batches;
 use lancedb::table::{
     AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration, NewColumnTransform,
@@ -16,6 +15,7 @@ use crate::error::NapiErrorExt;
 use crate::index::Index;
 use crate::merge::NativeMergeInsertBuilder;
 use crate::query::{Query, TakeQuery, VectorQuery};
+use crate::util::schema_to_buffer;
 
 #[napi]
 pub struct Table {
@@ -64,14 +64,7 @@ impl Table {
     #[napi(catch_unwind)]
     pub async fn schema(&self) -> napi::Result<Buffer> {
         let schema = self.inner_ref()?.schema().await.default_error()?;
-        let mut writer = FileWriter::try_new(vec![], &schema)
-            .map_err(|e| napi::Error::from_reason(format!("Failed to create IPC file: {}", e)))?;
-        writer
-            .finish()
-            .map_err(|e| napi::Error::from_reason(format!("Failed to finish IPC file: {}", e)))?;
-        Ok(Buffer::from(writer.into_inner().map_err(|e| {
-            napi::Error::from_reason(format!("Failed to get IPC file: {}", e))
-        })?))
+        schema_to_buffer(&schema)
     }
 
     #[napi(catch_unwind)]

nodejs/src/util.rs

@@ -1,7 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
+use arrow_ipc::writer::FileWriter;
+use arrow_schema::Schema;
 use lancedb::DistanceType;
+use napi::bindgen_prelude::Buffer;
 
 pub fn parse_distance_type(distance_type: impl AsRef<str>) -> napi::Result<DistanceType> {
     match distance_type.as_ref().to_lowercase().as_str() {
@@ -15,3 +18,15 @@ pub fn parse_distance_type(distance_type: impl AsRef<str>) -> napi::Result<Dista
         ))),
     }
 }
+
+/// Convert an Arrow Schema to an Arrow IPC file buffer
+pub fn schema_to_buffer(schema: &Schema) -> napi::Result<Buffer> {
+    let mut writer = FileWriter::try_new(vec![], schema)
+        .map_err(|e| napi::Error::from_reason(format!("Failed to create IPC file: {}", e)))?;
+    writer
+        .finish()
+        .map_err(|e| napi::Error::from_reason(format!("Failed to finish IPC file: {}", e)))?;
+    Ok(Buffer::from(writer.into_inner().map_err(|e| {
+        napi::Error::from_reason(format!("Failed to get IPC file: {}", e))
+    })?))
+}