feat: add timeout to query execution options (#2288)

Closes #2287


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added configurable timeout support for query executions. Users can now
specify maximum wait times for queries, enhancing control over
long-running operations across various integrations.
- **Tests**
- Expanded test coverage to validate timeout behavior in both
synchronous and asynchronous query flows, ensuring timely error
responses when query execution exceeds the specified limit.
- Introduced a new test suite to verify query operations when a timeout
is reached, checking for appropriate error handling.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

nodejs/__test__/table.test.ts

@@ -867,6 +867,44 @@ describe("When creating an index", () => {
   });
 });
 
+describe("When querying a table", () => {
+  let tmpDir: tmp.DirResult;
+  beforeEach(() => {
+    tmpDir = tmp.dirSync({ unsafeCleanup: true });
+  });
+  afterEach(() => tmpDir.removeCallback());
+
+  it("should throw an error when timeout is reached", async () => {
+    const db = await connect(tmpDir.name);
+    const data = makeArrowTable([
+      { text: "a", vector: [0.1, 0.2] },
+      { text: "b", vector: [0.3, 0.4] },
+    ]);
+    const table = await db.createTable("test", data);
+    await table.createIndex("text", { config: Index.fts() });
+
+    await expect(
+      table.query().where("text != 'a'").toArray({ timeoutMs: 0 }),
+    ).rejects.toThrow("Query timeout");
+
+    await expect(
+      table.query().nearestTo([0.0, 0.0]).toArrow({ timeoutMs: 0 }),
+    ).rejects.toThrow("Query timeout");
+
+    await expect(
+      table.search("a", "fts").toArray({ timeoutMs: 0 }),
+    ).rejects.toThrow("Query timeout");
+
+    await expect(
+      table
+        .query()
+        .nearestToText("a")
+        .nearestTo([0.0, 0.0])
+        .toArrow({ timeoutMs: 0 }),
+    ).rejects.toThrow("Query timeout");
+  });
+});
+
 describe("Read consistency interval", () => {
   let tmpDir: tmp.DirResult;
   beforeEach(() => {

nodejs/lancedb/query.ts

@@ -63,7 +63,7 @@ class RecordBatchIterable<
   // biome-ignore lint/suspicious/noExplicitAny: skip
   [Symbol.asyncIterator](): AsyncIterator<RecordBatch<any>, any, undefined> {
     return new RecordBatchIterator(
-      this.inner.execute(this.options?.maxBatchLength),
+      this.inner.execute(this.options?.maxBatchLength, this.options?.timeoutMs),
     );
   }
 }
@@ -79,6 +79,11 @@ export interface QueryExecutionOptions {
    * in smaller chunks.
    */
   maxBatchLength?: number;
+
+  /**
+   * Timeout for query execution in milliseconds
+   */
+  timeoutMs?: number;
 }
 
 /**
@@ -283,9 +288,11 @@ export class QueryBase<NativeQueryType extends NativeQuery | NativeVectorQuery>
     options?: Partial<QueryExecutionOptions>,
   ): Promise<NativeBatchIterator> {
     if (this.inner instanceof Promise) {
-      return this.inner.then((inner) => inner.execute(options?.maxBatchLength));
+      return this.inner.then((inner) =>
+        inner.execute(options?.maxBatchLength, options?.timeoutMs),
+      );
     } else {
-      return this.inner.execute(options?.maxBatchLength);
+      return this.inner.execute(options?.maxBatchLength, options?.timeoutMs);
     }
   }
 

nodejs/src/query.rs

@@ -131,11 +131,15 @@ impl Query {
     pub async fn execute(
         &self,
         max_batch_length: Option<u32>,
+        timeout_ms: Option<u32>,
     ) -> napi::Result<RecordBatchIterator> {
         let mut execution_opts = QueryExecutionOptions::default();
         if let Some(max_batch_length) = max_batch_length {
             execution_opts.max_batch_length = max_batch_length;
         }
+        if let Some(timeout_ms) = timeout_ms {
+            execution_opts.timeout = Some(std::time::Duration::from_millis(timeout_ms as u64))
+        }
         let inner_stream = self
             .inner
             .execute_with_options(execution_opts)
@@ -330,11 +334,15 @@ impl VectorQuery {
     pub async fn execute(
         &self,
         max_batch_length: Option<u32>,
+        timeout_ms: Option<u32>,
     ) -> napi::Result<RecordBatchIterator> {
         let mut execution_opts = QueryExecutionOptions::default();
         if let Some(max_batch_length) = max_batch_length {
             execution_opts.max_batch_length = max_batch_length;
         }
+        if let Some(timeout_ms) = timeout_ms {
+            execution_opts.timeout = Some(std::time::Duration::from_millis(timeout_ms as u64))
+        }
         let inner_stream = self
             .inner
             .execute_with_options(execution_opts)