mirror of
https://github.com/lancedb/lancedb.git
synced 2026-01-13 07:12:57 +00:00
8f8e06a2da
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.
33 lines
1.2 KiB
Rust
33 lines
1.2 KiB
Rust
// 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() {
|
|
"l2" => Ok(DistanceType::L2),
|
|
"cosine" => Ok(DistanceType::Cosine),
|
|
"dot" => Ok(DistanceType::Dot),
|
|
"hamming" => Ok(DistanceType::Hamming),
|
|
_ => Err(napi::Error::from_reason(format!(
|
|
"Invalid distance type '{}'. Must be one of l2, cosine, dot, or hamming",
|
|
distance_type.as_ref()
|
|
))),
|
|
}
|
|
}
|
|
|
|
/// 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))
|
|
})?))
|
|
}
|