Skip to content

Commit

Permalink
feat: implement EqDelFuture
Browse files Browse the repository at this point in the history
  • Loading branch information
sdd committed Feb 23, 2025
1 parent 8e90bdd commit ef827f1
Showing 1 changed file with 28 additions and 4 deletions.
32 changes: 28 additions & 4 deletions crates/iceberg/src/arrow/delete_file_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,11 @@ use std::collections::HashMap;
use std::future::Future;
use std::ops::BitAndAssign;
use std::pin::Pin;
use std::sync::{Arc, RwLock};
use std::sync::{Arc, OnceLock, RwLock};
use std::task::{Context, Poll};

use futures::channel::oneshot;
use futures::future::join_all;
use futures::{StreamExt, TryStreamExt};
use roaring::RoaringTreemap;

Expand All @@ -41,19 +42,32 @@ use crate::{Error, ErrorKind, Result};
// state that can be awaited upon to get te EQ deletes. That way we can check to see if
// a load of each Eq delete file is already in progress and avoid starting another one.
#[derive(Debug, Clone)]
struct EqDelFuture {}
struct EqDelFuture {
result: OnceLock<Predicate>,
}

impl EqDelFuture {
pub fn new() -> (oneshot::Sender<Predicate>, Self) {
todo!()
let (tx, rx) = oneshot::channel();
let result = OnceLock::new();

crate::runtime::spawn({
let result = result.clone();
async move { result.set(rx.await.unwrap()) }
});

(tx, Self { result })
}
}

impl Future for EqDelFuture {
type Output = Predicate;

fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
todo!()
match self.result.get() {
None => Poll::Pending,
Some(predicate) => Poll::Ready(predicate.clone()),
}
}
}

Expand Down Expand Up @@ -189,6 +203,16 @@ impl DeleteFileManager {
.try_collect::<Vec<_>>()
.await?;

// wait for all in-progress EQ deletes from other tasks
let _ = join_all(results.iter().filter_map(|i| {
if let ParsedDeleteFileContext::InProgEqDel(fut) = i {
Some(fut.clone())
} else {
None
}
}))
.await;

let merged_delete_vectors = results
.into_iter()
.fold(HashMap::default(), Self::merge_delete_vectors);
Expand Down

0 comments on commit ef827f1

Please sign in to comment.