lancedb 0.27.1

LanceDB: A serverless, low-latency vector database for AI applications
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors

use arrow_array::RecordBatch;
use futures::TryStreamExt;
use lance_datagen::{BatchCount, BatchGeneratorBuilder, RowCount};

use crate::{
    Error, Table,
    arrow::{SendableRecordBatchStream, SimpleRecordBatchStream},
    connect,
};

#[async_trait::async_trait]
pub trait LanceDbDatagenExt {
    async fn into_mem_table(
        self,
        table_name: &str,
        rows_per_batch: RowCount,
        num_batches: BatchCount,
    ) -> Table;
}

#[async_trait::async_trait]
impl LanceDbDatagenExt for BatchGeneratorBuilder {
    async fn into_mem_table(
        self,
        table_name: &str,
        rows_per_batch: RowCount,
        num_batches: BatchCount,
    ) -> Table {
        let (stream, schema) = self.into_reader_stream(rows_per_batch, num_batches);
        let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream::new(
            stream.map_err(Error::from),
            schema,
        ));
        let db = connect("memory:///").execute().await.unwrap();
        db.create_table(table_name, stream).execute().await.unwrap()
    }
}

pub async fn virtual_table(name: &str, values: &RecordBatch) -> Table {
    let schema = values.schema();
    let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream::new(
        futures::stream::once(std::future::ready(Ok(values.clone()))),
        schema,
    ));
    let db = connect("memory:///").execute().await.unwrap();
    db.create_table(name, stream).execute().await.unwrap()
}