// Copyright 2023 LanceDB Developers. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. import { Vector, tableFromIPC } from 'apache-arrow' import { type EmbeddingFunction } from './embedding/embedding_function' import { type MetricType } from '.' // eslint-disable-next-line @typescript-eslint/no-var-requires const { tableSearch } = require('../native.js') /** * A builder for nearest neighbor queries for LanceDB. */ export class Query { private readonly _query?: T private readonly _tbl?: any private _queryVector?: number[] private _limit?: number private _refineFactor?: number private _nprobes: number private _select?: string[] private _filter?: string private _metricType?: MetricType private _prefilter: boolean private _fastSearch: boolean protected readonly _embeddings?: EmbeddingFunction constructor (query?: T, tbl?: any, embeddings?: EmbeddingFunction) { this._tbl = tbl this._query = query this._limit = 10 this._nprobes = 20 this._refineFactor = undefined this._select = undefined this._filter = undefined this._metricType = undefined this._embeddings = embeddings this._prefilter = false this._fastSearch = false } /*** * Sets the number of results that will be returned * default value is 10 * @param value number of results */ limit (value: number): Query { this._limit = value return this } /** * Refine the results by reading extra elements and re-ranking them in memory. * @param value refine factor to use in this query. */ refineFactor (value: number): Query { this._refineFactor = value return this } /** * The number of probes used. A higher number makes search more accurate but also slower. * @param value The number of probes used. */ nprobes (value: number): Query { this._nprobes = value return this } /** * A filter statement to be applied to this query. * @param value A filter in the same format used by a sql WHERE clause. */ filter (value: string): Query { this._filter = value return this } where = this.filter /** Return only the specified columns. * * @param value Only select the specified columns. If not specified, all columns will be returned. */ select (value: string[]): Query { this._select = value return this } /** * The MetricType used for this Query. * @param value The metric to the. @see MetricType for the different options */ metricType (value: MetricType): Query { this._metricType = value return this } prefilter (value: boolean): Query { this._prefilter = value return this } /** * Skip searching un-indexed data. This can make search faster, but will miss * any data that is not yet indexed. */ fastSearch (value: boolean): Query { this._fastSearch = value return this } /** * Execute the query and return the results as an Array of Objects */ async execute> (): Promise { if (this._query !== undefined) { if (this._embeddings !== undefined) { this._queryVector = (await this._embeddings.embed([this._query]))[0] } else { this._queryVector = this._query as number[] } } const isElectron = this.isElectron() const buffer = await tableSearch.call(this._tbl, this, isElectron) const data = tableFromIPC(buffer) return data.toArray().map((entry: Record) => { const newObject: Record = {} Object.keys(entry).forEach((key: string) => { if (entry[key] instanceof Vector) { // toJSON() returns f16 array correctly newObject[key] = (entry[key] as any).toJSON() } else { newObject[key] = entry[key] as any } }) return newObject as unknown as T }) } // See https://github.com/electron/electron/issues/2288 private isElectron (): boolean { try { // eslint-disable-next-line no-prototype-builtins return (process?.versions?.hasOwnProperty('electron') || navigator?.userAgent?.toLowerCase()?.includes(' electron')) } catch (e) { return false } } }