34 posts tagged with "Rust"

use anyhow::Result;
use arrow::array::AsArray;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::fs::File;

const PQ_FILE: &str = "../assets/sample.parquet";

fn main() -> Result<()> {
  read_with_parquet(PQ_FILE)?;
  Ok(())
}

fn read_with_parquet(file: &str) -> Result<()> {
  let file = File::open(file)?;
  let reader = ParquetRecordBatchReaderBuilder::try_new(file)?
    .with_batch_size(8192)
    .with_limit(3)
    .build()?;

  for record_batch in reader {
    let record_batch = record_batch?;
    let emails = record_batch.column(0).as_binary::<i32>();

    for email in emails {
      let email = email.unwrap();
      println!("{:?}", String::from_utf8_lossy(email));
    }
  }
  Ok(())
}

use datafusion::prelude::*;
use datafusion::error::Result;

let ctx = SessionContext::new();
ctx.register_csv("example", "assets/example.csv", CsvReadOption::new()).await?;

pub async fn register_csv(
	&self,
	name: &str,
	table_path: &str,
	options: CsvReadOptions<'_>,
) -> Result<()> {
	let listing_options = options
		.to_listing_options(&self.copied_config(), self.copied_table_options());

	self.register_listing_table(
		name,
		table_path,
		listing_options,
		options.schema.map(|s| Arc::new(s.to_owned())),
		None,
	)
	.await?;

	Ok(())
}

let df = ctx.sql("SELECT a, MIN(b), FROM example WHERE a <= b GROUP BY a LIMIT 100").await?;

pub async fn sql(&self, sql: &str) -> Result<DataFrame> {
  self.sql_with_options(sql, SQLOptions::new()).await
}

df.show().await?;

pub async fn show(self) -> Result<()> {
	let results = self.collect().await?;
	Ok(pretty::print_batches(&results)?)
}

use arrow_schema::{DataType, Field, Schema};
use lancedb::{connect, Result};
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<()> {
  let created_empty_talbe = create_empty_table().await?;
  println!(
    "Created empty table: {}, Table name: {}",
    created_empty_talbe,
    created_empty_talbe.name()
  );
  Ok(())
}

#[allow(unused)]
async fn create_empty_table() -> Result<LanceDbTable> {
	// 创建模式定义
  let schema = Arc::new(Schema::new(vec![
    Field::new("id", DataType::Int32, false),
    Field::new("name", DataType::Utf8, false),
  ]));
	// 创建数据库URI目录
  let uri = "data/sample-lancedb";
	// 连接数据库
  let db = connect(uri).execute().await?;
	// 创建一个空表
  let table = db
    .create_empty_table("empty_talbe", schema)
    .execute()
    .await?;
  Ok(table)
}

use arrow_schema::{DataType, Field, Schema};
use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray};
use lancedb::{connect, Result, Table as LanceDbTable};
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<()> {
  let created_table_with_data = create_table_with_data().await?;
  println!(
    "Created table with data: {}, Table name: {}",
    created_table_with_data,
    created_table_with_data.name()
  );
  Ok(())
}

#[allow(unused)]
async fn create_table_with_data() -> Result<LanceDbTable> {
  // 创建本地数据库URI目录
  let uri = "data/sample-lancedb";
  // 连接数据库
  let db = connect(uri).execute().await?;

  // 创建模式定义
  let schema = Arc::new(Schema::new(vec![
    Field::new("id", DataType::Int32, false),
    Field::new("name", DataType::Utf8, false),
  ]));

  // 初始化`ids`列的数据
  let ids = Int32Array::from(vec![1, 2, 3]); 
  // 初始化`name`列的数据
  let names = StringArray::from(vec!["Alice", "Bob", "Lily"]); 
  // 使用`Schema`以及列数据创建`RecordBatch`
  let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(names)])?;
  // 使用`RecordBatch`创建`RecordBatchIterator`
  let batchs = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema);
  // 创建表，并插入初始化数据
  let table = db
    .create_table("table_with_person", batchs)
    .execute()
    .await?;
  Ok(table)
}

use arrow_array::types::Float32Type;
use arrow_array::{FixedSizeListArray, Int32Array, RecordBatch, RecordBatchIterator, StringArray};
use arrow_schema::{DataType, Field, Schema};
use lancedb::arrow::IntoArrow;
use lancedb::{connect, Result, Table as LanceDbTable};
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<()> {
  let created_table_with_records = create_table_with_records().await?;
  println!(
    "Created table with records: {}, Table name: {}",
    created_table_with_records,
    created_table_with_records.name()
  );
  Ok(())
}

#[allow(unused)]
async fn create_table_with_records() -> Result<LanceDbTable> {
  let uri = "data/sample-lancedb";
  let db = connect(uri).execute().await?;

  let initial_data = create_some_records()?;
  let tbl = db.create_table("my_table", initial_data).execute().await?;

  let new_data = create_some_records()?;
  // NOTICE: 只有实现了 IntoArrow 的类型才能使用`add`方法，即`create_some_records`返回的类型
  tbl.add(new_data).execute().await?;
  Ok(tbl)
}

#[allow(unused)]
fn create_some_records() -> Result<impl IntoArrow> {
  const TOTAL: usize = 1000;
  const DIM: usize = 128;

  let schema = Arc::new(Schema::new(vec![
    Field::new("id", DataType::Int32, false),
    Field::new(
      "vector",
      DataType::FixedSizeList(
        Arc::new(Field::new("item", DataType::Float32, true)),
        DIM as i32,
      ),
      true,
    ),
  ]));

  let batch = RecordBatch::try_new(
    schema.clone(),
    vec![
      Arc::new(Int32Array::from_iter_values(0..TOTAL as i32)),
      Arc::new(
        FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
          (0..TOTAL).map(|_| Some(vec![Some(1.0); DIM])),
          DIM as i32,
        ),
      ),
    ],
  )
  .unwrap();
  let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone());
  Ok(Box::new(batches))
}

use lancedb::{connect, Result, Table as LanceDbTable};

async fn main() -> Result<()> {
  let opened_table = open_with_existing_table().await?;
  println!(
    "Opened table: {}, Table name: {}",
    opened_table,
    opened_table.name()
  );
}

#[allow(unused)]
async fn open_with_existing_table() -> Result<LanceDbTable> {
  let uri = "data/sample-lancedb";
  let db = connect(uri).execute().await?;
  let table = db.open_table("my_table").execute().await?;
  Ok(table)
}

use lancedb::{connect, Result};

#[tokio::main]
async fn main() -> Result<()> {
  let queried_result = query_table().await?;
  println!("Query result: {:?}", queried_result);

  delete_table_records().await?; // 根据条件删除表中的记录
  Ok(())
}

#[allow(unused)]
async fn delete_table_records() -> Result<()> {
  let uri = "data/sample-lancedb";
  let db = connect(uri).execute().await?;
  let table = db.open_table("my_table").execute().await?;
  table.delete("id > 24").await?;
  Ok(())
}

use lancedb::{connect, Result};

#[tokio::main]
async fn main() -> Result<()> {
  let queried_result = query_table().await?;
  println!("Query result: {:?}", queried_result);

  drop_table().await?; // 删除 data/sample-lancedb/my_table
  Ok(())
}

#[allow(unused)]
async fn drop_table() -> Result<()> {
  let uri = "data/sample-lancedb";
  let db = connect(uri).execute().await?;
  db.drop_table("my_table").await?;
  Ok(())
}

use lancedb::{connect, Result};

#[tokio::main]
async fn main() -> Result<()> {
  let queried_result = query_table().await?;
  println!("Query result: {:?}", queried_result);

  drop_database().await?; // 删除 data/sample-lancedb
  Ok(())
}

#[allow(unused)]
async fn drop_database() -> Result<()> {
  let uri = "data/sample-lancedb";
  let db = connect(uri).execute().await?;
  db.drop_db().await?;
  Ok(())
}

use lancedb::{connect, Result};

#[tokio::main]
async fn main() -> Result<()> {

  let queried_result = query_table().await?;
  println!("Query result: {:?}", queried_result);
  Ok(())
}

#[allow(unused)]
async fn query_table() -> Result<VectorQuery> {
  let uri = "data/sample-lancedb";
  let db = connect(uri).execute().await?;
  let table = db.open_table("my_table").execute().await?;
  let result = table.query().nearest_to(&[1.0; 128])?;
  Ok(result)
}

use lancedb::{connect, Result};

#[tokio::main]
async fn main() -> Result<()> {

  update_table().await?;
  Ok(())
}

#[allow(unused)]
async fn update_table() -> Result<()> {
  let uri = "data/sample-lancedb";
  let db = connect(uri).execute().await?;
  let table = db.open_table("table_with_person").execute().await?;
  println!("Before update: {:?}", table.query());
  table
    .update()
    .only_if("id=0")
    .column("name", "Bob")
    .execute()
    .await?; // Bob -> Halzzz

  Ok(())
}

use lancedb::{connect, Result};

#[tokio::main]
async fn main() -> Result<()> {
  uri = "data/example-lancedb";
  db_builder = connect(uri);
  db_connect = connect(uri).execute().await?;
  println!("LanceDB builder: {:}", db_builder);
  println!("LanceDB connect: {}", db_connect);
  Ok(())
}

/// Connect to a LanceDB database.
///
/// # Arguments
///
/// * `uri` - URI where the database is located, can be a local directory, supported remote cloud storage,
///           or a LanceDB Cloud database.  See [ConnectOptions::uri] for a list of accepted formats
pub fn connect(uri: &str) -> ConnectBuilder {
  ConnectBuilder::new(uri)
}

pub struct ConnectBuilder {
  /// Database URI
  ///
  /// ### Accpeted URI formats
  ///
  /// - `/path/to/database` - local database on file system.
  /// - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud object store
  /// - `db://dbname` - LanceDB Cloud
  uri: String,

  /// LanceDB Cloud API key, required if using Lance Cloud
  api_key: Option<String>,
  /// LanceDB Cloud region, required if using Lance Cloud
  region: Option<String>,
  /// LanceDB Cloud host override, only required if using an on-premises Lance Cloud instance
  host_override: Option<String>,

  storage_options: HashMap<String, String>,

  /// The interval at which to check for updates from other processes.
  ///
  /// If None, then consistency is not checked. For performance
  /// reasons, this is the default. For strong consistency, set this to
  /// zero seconds. Then every read will check for updates from other
  /// processes. As a compromise, you can set this to a non-zero timedelta
  /// for eventual consistency. If more than that interval has passed since
  /// the last check, then the table will be checked for updates. Note: this
  /// consistency only applies to read operations. Write operations are
  /// always consistent.
  read_consistency_interval: Option<std::time::Duration>,
  embedding_registry: Option<Arc<dyn EmbeddingRegistry>>,
}

impl ConnectBuilder {
  /// Create a new [`ConnectOptions`] with the given database URI.
  pub fn new(uri: &str) -> Self {
    Self {
      uri: uri.to_string(),
      api_key: None,
      region: None,
      host_override: None,
      read_consistency_interval: None,
      storage_options: HashMap::new(),
      embedding_registry: None,
    }
  }

  // ......

  /// Establishes a connection to the database
  pub async fn execute(self) -> Result<Connection> {
    if self.uri.starts_with("db") {
      self.execute_remote()
    } else {
      let internal = Arc::new(Database::connect_with_options(&self).await?);
      Ok(Connection {
           internal,
           uri: self.uri,
      })
    }
  }
}

#[derive(Debug)]
struct Database {
  object_store: ObjectStore,
  query_string: Option<String>,

  pub(crate) uri: String,
  pub(crate) base_path: object_store::path::Path,

  // the object store wrapper to use on write path
  pub(crate) store_wrapper: Option<Arc<dyn WrappingObjectStore>>,

  read_consistency_interval: Option<std::time::Duration>,

  // Storage options to be inherited by tables created from this connection
  storage_options: HashMap<String, String>,
  embedding_registry: Arc<dyn EmbeddingRegistry>,
}

impl Database {
  async fn connect_with_options(options: &ConnectBuilder) -> Result<Self> {
    let uri = &options.uri;
    let parse_res = url::Url::parse(uri);

    // TODO: pass params regardless of OS
    match parse_res {
      Ok(url) if url.scheme().len() == 1 && cfg!(windows) => {
        Self::open_path(
          uri,
          options.read_consistency_interval,
          options.embedding_registry.clone(),
        )
       .await
     }
     Ok(mut url) => {
       // iter thru the query params and extract the commit store param
       let mut engine = None;
       let mut mirrored_store = None;
       let mut filtered_querys = vec![];

       // WARNING: specifying engine is NOT a publicly supported feature in lancedb yet
       // THE API WILL CHANGE
       for (key, value) in url.query_pairs() {
         if key == ENGINE {
           engine = Some(value.to_string());
         } else if key == MIRRORED_STORE {
           if cfg!(windows) {
             return Err(Error::NotSupported {
               message: "mirrored store is not supported on windows".into(),
             });
           }
           mirrored_store = Some(value.to_string());
         } else {
           // to owned so we can modify the url
           filtered_querys.push((key.to_string(), value.to_string()));
         }
       }

       // Filter out the commit store query param -- it's a lancedb param
       url.query_pairs_mut().clear();
       url.query_pairs_mut().extend_pairs(filtered_querys);
       // Take a copy of the query string so we can propagate it to lance
       let query_string = url.query().map(|s| s.to_string());
       // clear the query string so we can use the url as the base uri
       // use .set_query(None) instead of .set_query("") because the latter
       // will add a trailing '?' to the url
       url.set_query(None);

       let table_base_uri = if let Some(store) = engine {
         static WARN_ONCE: std::sync::Once = std::sync::Once::new();
         WARN_ONCE.call_once(|| {
           log::warn!("Specifing engine is not a publicly supported feature in lancedb yet. THE API WILL CHANGE");
         });
         let old_scheme = url.scheme().to_string();
         let new_scheme = format!("{}+{}", old_scheme, store);
         url.to_string().replacen(&old_scheme, &new_scheme, 1)
       } else {
         url.to_string()
       };

       let plain_uri = url.to_string();

       let storage_options = options.storage_options.clone();
       let os_params = ObjectStoreParams {
         storage_options: Some(storage_options.clone()),
         ..Default::default()
       };
       let (object_store, base_path) =
         ObjectStore::from_uri_and_params(&plain_uri, &os_params).await?;
       if object_store.is_local() {
         Self::try_create_dir(&plain_uri).context(CreateDirSnafu { path: plain_uri })?;
       }

       let write_store_wrapper = match mirrored_store {
         Some(path) => {
           let mirrored_store = Arc::new(LocalFileSystem::new_with_prefix(path)?);
           let wrapper = MirroringObjectStoreWrapper::new(mirrored_store);
           Some(Arc::new(wrapper) as Arc<dyn WrappingObjectStore>)
         }
         None => None,
       };

       let embedding_registry = options
           .embedding_registry
           .clone()
           .unwrap_or_else(|| Arc::new(MemoryRegistry::new()));
       Ok(Self {
         uri: table_base_uri,
         query_string,
         base_path,
         object_store,
         store_wrapper: write_store_wrapper,
         read_consistency_interval: options.read_consistency_interval,
         storage_options,
         embedding_registry,
       })
      }
      Err(_) => {
        Self::open_path(
          uri,
          options.read_consistency_interval,
          options.embedding_registry.clone(),
        )
        .await
      }
    }
  }
}

/// A connection to LanceDB
#[derive(Clone)]
pub struct Connection {
  uri: String,
  internal: Arc<dyn ConnectionInternal>,
}

// request_id.rs
use super::REQUEST_ID_HEADER;
use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response};
use tracing::warn;

pub async fn set_request_id(mut req: Request, next: Next) -> Response {
  let id = match req.headers().get(REQUEST_ID_HEADER) {
    Some(v) => Some(v.clone()),
    None => {
      let request_id = uuid::Uuid::new_v4().to_string();
      match HeaderValue::from_str(&request_id) {
        Ok(v) => {
          req.headers_mut().insert(REQUEST_ID_HEADER, v.clone());
          Some(v)
        }
        Err(e) => {
          warn!("parse generated request id failed: {}", e);
          None
        }
      }
    }
  };

  let mut res = next.run(req).await;

  if let Some(id) = id {
    res.headers_mut().insert(REQUEST_ID_HEADER, id);
  }
  res
}

// auth.rs
use crate::AppState;
use axum::{
  extract::{FromRequestParts, Request, State},
  http::StatusCode,
  middleware::Next,
  response::{IntoResponse, Response},
};
use axum_extra::{
  headers::{authorization::Bearer, Authorization},
  TypedHeader;
use tracing::warn;

pub async fn verify_token(State(state): State<AppState>, req: Request, next: Next) -> Response {
  let (mut parts, body) = req.into_parts();
  let req = match TypedHeader::<Authorization<Bearer>>::from_request_parts(&mut parts, &state).await {
    Ok(TypedHeader(Authorization(bearer))) => {
      let token = bearer.token();
      match state.dk.verify(token) {
        Ok(user) => {
          let mut req = Request::from_parts(parts, body);
          req.extensions_mut().insert(user);
          req
        }
        Err(e) => {
          let msg = format!("verify token failed: {}", e);
          warn!(msg);
          return (StatusCode::FORBIDDEN, msg).into_response();
        }
      }
    }
    Err(e) => {
      let msg = format!("parse Authorization header failed: {}", e);
      warn!(msg);
      return (StatusCode::UNAUTHORIZED, msg).into_response();
    }
  };
  next.run(req).await
}

// server_time.rs
use super::{REQUEST_ID_HEADER, SERVER_TIME_HEADER};
use axum::{extract::Request, response::Response};
use std::{
  future::Future,
  pin::Pin,
  task::{Context, Poll},
};
use tokio::time::Instant;
use tower::{Layer, Service};
use tracing::warn;

#[derive(Clone)]
pub struct ServerTimeLayer;

impl<S> Layer<S> for ServerTimeLayer {
  type Service = ServerTimeMiddleware<S>;

  fn layer(&self, inner: S) -> Self::Service {
    ServerTimeMiddleware { inner }
  }
}

#[derive(Clone)]
pub struct ServerTimeMiddleware<S> {
  inner: S,
}

impl<S> Service<Request> for ServerTimeMiddleware<S>
where
  S: Service<Request, Response = Response> + Send + 'static,
  S::Future: Send + 'static,
{
  type Response = S::Response;
  type Error = S::Error;
  type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
    self.inner.poll_ready(cx)
  }

  fn call(&mut self, req: Request) -> Self::Future {
    let start = Instant::now();
    let future = self.inner.call(req);
    Box::pin(async move {
      let mut res: Response = future.await?;
      let elapsed = format!("{}us", start.elapsed().as_micros());
      match elapsed.parse() {
        Ok(v) => {
          res.headers_mut().insert(SERVER_TIME_HEADER, v);
        }
        Err(e) => {
          warn!(
            "Parse elapsed time failed: {} for request {:?}",
            e,
            res.headers().get(REQUEST_ID_HEADER)
          );
        }
      }
      Ok(res)
    })
  }
}