remove float[] array support, only use vector
This commit is contained in:
@@ -1,5 +1,4 @@
|
||||
use anyhow::{Context, Result};
|
||||
use bytes::BytesMut;
|
||||
use clap::Parser;
|
||||
use deadpool_postgres::{Config, Pool, Runtime};
|
||||
use dotenv::dotenv;
|
||||
@@ -57,8 +56,6 @@ async fn get_column_types(
|
||||
let pg_type = match (data_type, udt_name) {
|
||||
("integer", _) => Type::INT4,
|
||||
("bigint", _) => Type::INT8,
|
||||
("ARRAY", "_float4") => Type::FLOAT4_ARRAY,
|
||||
("ARRAY", "_float8") => Type::FLOAT8_ARRAY,
|
||||
("USER-DEFINED", type_name) => pg_types::get_pg_type(client, type_name).await?,
|
||||
_ => anyhow::bail!("Unsupported column type: {} ({})", data_type, udt_name),
|
||||
};
|
||||
@@ -152,11 +149,14 @@ async fn save_embeddings(
|
||||
&[&args.target_id_column, &args.target_embedding_column],
|
||||
)
|
||||
.await?;
|
||||
info!("Target table column types: {:?}", types);
|
||||
|
||||
if types.len() != 2 {
|
||||
anyhow::bail!("Failed to get both column types from target table");
|
||||
}
|
||||
|
||||
// Create a temporary table with the same structure
|
||||
info!("Creating temporary table...");
|
||||
let id_type = match types[0] {
|
||||
let id_type_str = match types[0] {
|
||||
Type::INT4 => "INTEGER",
|
||||
Type::INT8 => "BIGINT",
|
||||
_ => anyhow::bail!("Unexpected type for ID column: {:?}", types[0]),
|
||||
@@ -166,7 +166,7 @@ async fn save_embeddings(
|
||||
{} {},
|
||||
{} vector({})
|
||||
)",
|
||||
args.target_id_column, id_type, args.target_embedding_column, args.factors
|
||||
args.target_id_column, id_type_str, args.target_embedding_column, args.factors
|
||||
);
|
||||
info!("Temp table creation SQL: {}", create_temp);
|
||||
client.execute(&create_temp, &[]).await?;
|
||||
@@ -327,6 +327,13 @@ async fn validate_schema(client: &deadpool_postgres::Client, args: &Args) -> Res
|
||||
let udt_name: &str = row.get("udt_name");
|
||||
if udt_name == "vector" {
|
||||
let vector_dim: i32 = row.get("vector_dim");
|
||||
if vector_dim <= 0 {
|
||||
anyhow::bail!(
|
||||
"Invalid vector dimension {} for column '{}'",
|
||||
vector_dim,
|
||||
args.target_embedding_column,
|
||||
);
|
||||
}
|
||||
if vector_dim != args.factors as i32 {
|
||||
anyhow::bail!(
|
||||
"Vector dimension mismatch: column '{}' has dimension {}, but factors is {}",
|
||||
|
||||
Reference in New Issue
Block a user