Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Silas fix streaming bug#1406

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Merged
SilasMarvin merged 4 commits intomasterfromsilas-fix-streaming-bug
Apr 15, 2024
Merged
Show file tree
Hide file tree
Changes fromall commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletionpgml-sdks/pgml/Cargo.toml
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
[package]
name ="pgml"
version ="1.0.1"
version ="1.0.2"
edition ="2021"
authors = ["PosgresML <team@postgresml.org>"]
homepage ="https://postgresml.org/"
Expand Down
2 changes: 1 addition & 1 deletionpgml-sdks/pgml/javascript/package.json
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
{
"name":"pgml",
"version":"1.0.1",
"version":"1.0.2",
"description":"Open Source Alternative for Building End-to-End Vector Search Applications without OpenAI & Pinecone",
"keywords": [
"postgres",
Expand Down
2 changes: 1 addition & 1 deletionpgml-sdks/pgml/pyproject.toml
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -5,7 +5,7 @@ build-backend = "maturin"
[project]
name ="pgml"
requires-python =">=3.7"
version ="1.0.1"
version ="1.0.2"
description ="Python SDK is designed to facilitate the development of scalable vector search applications on PostgreSQL databases."
authors = [
{name ="PostgresML",email ="team@postgresml.org"},
Expand Down
140 changes: 34 additions & 106 deletionspgml-sdks/pgml/src/transformer_pipeline.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,6 @@
use anyhow::Context;
use futures::Stream;
use rust_bridge::{alias, alias_methods};
use sqlx::{postgres::PgRow, Row};
use sqlx::{Postgres, Transaction};
use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::task::Poll;
use sqlx::Row;
use tracing::instrument;

/// Provides access to builtin database methods
Expand All@@ -22,99 +16,6 @@ use crate::{get_or_initialize_pool, types::Json};
#[cfg(feature = "python")]
use crate::types::{GeneralJsonAsyncIteratorPython, JsonPython};

#[allow(clippy::type_complexity)]
struct TransformerStream {
transaction: Option<Transaction<'static, Postgres>>,
future: Option<Pin<Box<dyn Future<Output = Result<Vec<PgRow>, sqlx::Error>> + Send + 'static>>>,
commit: Option<Pin<Box<dyn Future<Output = Result<(), sqlx::Error>> + Send + 'static>>>,
done: bool,
query: String,
db_batch_size: i32,
results: VecDeque<PgRow>,
}

impl std::fmt::Debug for TransformerStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TransformerStream").finish()
}
}

impl TransformerStream {
fn new(transaction: Transaction<'static, Postgres>, db_batch_size: i32) -> Self {
let query = format!("FETCH {} FROM c", db_batch_size);
Self {
transaction: Some(transaction),
future: None,
commit: None,
done: false,
query,
db_batch_size,
results: VecDeque::new(),
}
}
}

impl Stream for TransformerStream {
type Item = anyhow::Result<Json>;

fn poll_next(
mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
if self.done {
if let Some(c) = self.commit.as_mut() {
if c.as_mut().poll(cx).is_ready() {
self.commit = None;
}
}
} else {
if self.future.is_none() {
unsafe {
let s = self.as_mut().get_unchecked_mut();
let s: *mut Self = s;
let s = Box::leak(Box::from_raw(s));
s.future = Some(Box::pin(
sqlx::query(&s.query).fetch_all(&mut **s.transaction.as_mut().unwrap()),
));
}
}

if let Poll::Ready(o) = self.as_mut().future.as_mut().unwrap().as_mut().poll(cx) {
let rows = o?;
if rows.len() < self.db_batch_size as usize {
self.done = true;
unsafe {
let s = self.as_mut().get_unchecked_mut();
let transaction = std::mem::take(&mut s.transaction).unwrap();
s.commit = Some(Box::pin(transaction.commit()));
}
} else {
unsafe {
let s = self.as_mut().get_unchecked_mut();
let s: *mut Self = s;
let s = Box::leak(Box::from_raw(s));
s.future = Some(Box::pin(
sqlx::query(&s.query).fetch_all(&mut **s.transaction.as_mut().unwrap()),
));
}
}
for r in rows.into_iter() {
self.results.push_back(r)
}
}
}

if !self.results.is_empty() {
let r = self.results.pop_front().unwrap();
Poll::Ready(Some(Ok(r.get::<Json, _>(0))))
} else if self.done {
Poll::Ready(None)
} else {
Poll::Pending
}
}
}

#[alias_methods(new, transform, transform_stream)]
impl TransformerPipeline {
/// Creates a new [TransformerPipeline]
Expand DownExpand Up@@ -200,7 +101,7 @@ impl TransformerPipeline {
) -> anyhow::Result<GeneralJsonAsyncIterator> {
let pool = get_or_initialize_pool(&self.database_url).await?;
let args = args.unwrap_or_default();
let batch_size = batch_size.unwrap_or(10);
let batch_size = batch_size.unwrap_or(1);

let mut transaction = pool.begin().await?;
// We set the task in the new constructor so we can unwrap here
Expand DownExpand Up@@ -234,10 +135,37 @@ impl TransformerPipeline {
.await?;
}

Ok(GeneralJsonAsyncIterator(Box::pin(TransformerStream::new(
transaction,
batch_size,
))))
let s = futures::stream::try_unfold(transaction, move |mut transaction| async move {
let query = format!("FETCH {} FROM c", batch_size);
let mut res: Vec<Json> = sqlx::query_scalar(&query)
.fetch_all(&mut *transaction)
.await?;
if !res.is_empty() {
if batch_size > 1 {
let res: Vec<String> = res
.into_iter()
.map(|v| {
v.0.as_array()
.context("internal SDK error - cannot parse db value as array. Please post a new github issue")
.map(|v| {
v[0].as_str()
.context(
"internal SDK error - cannot parse db value as string. Please post a new github issue",
)
.map(|v| v.to_owned())
})
})
.collect::<anyhow::Result<anyhow::Result<Vec<String>>>>()??;
Ok(Some((serde_json::json!(res).into(), transaction)))
} else {
Ok(Some((std::mem::take(&mut res[0]), transaction)))
}
} else {
transaction.commit().await?;
Ok(None)
}
});
Ok(GeneralJsonAsyncIterator(Box::pin(s)))
}
}

Expand DownExpand Up@@ -305,7 +233,7 @@ mod tests {
serde_json::json!("AI is going to").into(),
Some(
serde_json::json!({
"max_new_tokens":10
"max_new_tokens":30
})
.into(),
),
Expand Down
8 changes: 3 additions & 5 deletionspgml-sdks/pgml/src/types.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
use anyhow::Context;
use futures::{Stream, StreamExt};
use futures::{stream::BoxStream,Stream, StreamExt};
use itertools::Itertools;
use rust_bridge::alias_manual;
use sea_query::Iden;
Expand DownExpand Up@@ -123,11 +123,9 @@ impl IntoTableNameAndSchema for String {
}
}

/// A wrapper around `std::pin::Pin<Box<dyn Stream<Item =anyhow::Result<Json>> + Send>>`
/// A wrapper around `BoxStream<'static,anyhow::Result<Json>>`
#[derive(alias_manual)]
pub struct GeneralJsonAsyncIterator(
pub std::pin::Pin<Box<dyn Stream<Item = anyhow::Result<Json>> + Send>>,
);
pub struct GeneralJsonAsyncIterator(pub BoxStream<'static, anyhow::Result<Json>>);

impl Stream for GeneralJsonAsyncIterator {
type Item = anyhow::Result<Json>;
Expand Down
[8]ページ先頭
©2009-2025 Movatter.jp