Skip to content

DRAFT [SS-348] Iceberg sink: Write data files straight from the arrangement batches. #37686

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.

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Closed
ublubu wants to merge 2 commits into from
Closed

DRAFT [SS-348] Iceberg sink: Write data files straight from the arrangement batches. #37686

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0ed3911
batch diff_pairs iterator, preparing for an async traversal
ublubu Jul 14, 2026
e556df4
Write data files straight from the arrangement batches.
ublubu Jul 14, 2026
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
311 changes: 261 additions & 50 deletions src/interchange/src/envelopes.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,80 +19,154 @@ use mz_ore::cast::CastFrom;
use mz_repr::{
CatalogItemId, ColumnName, Datum, Diff, Row, RowPacker, SqlColumnType, SqlScalarType,
};
use timely::progress::Antichain;

use crate::avro::DiffPair;

/// Walks `batch` and invokes `on_diff_pair` for each `DiffPair` at each
/// `(key, timestamp)`.
///
/// Thin wrapper around `iter_diff_pairs`.
pub fn for_each_diff_pair<B, C, F>(batch: &B, mut on_diff_pair: F)
where
B: BatchReader<Time = C::Time> + Navigable<Cursor = C>,
C: Cursor<Storage = B, Diff = Diff>,
C::Time: Copy,
C::ValOwn: 'static,
F: FnMut(&<C::KeyContainer as BatchContainer>::Owned, C::Time, DiffPair<C::ValOwn>),
{
for (key, timed_pairs) in iter_diff_pairs(batch, None, None) {
for (time, pair) in timed_pairs {
on_diff_pair(&key, time, pair);
}
}
}

/// Just like `for_each_diff_pair`, but async and fallible.
///
/// Thin wrapper around `iter_diff_pairs`.
pub async fn for_each_diff_pair_async<B, C, F, E>(
batch: &B,
lower: Option<Antichain<C::Time>>,
upper: Option<Antichain<C::Time>>,
mut on_diff_pair: F,
) -> Result<(), E>
where
B: BatchReader<Time = C::Time> + Navigable<Cursor = C>,
C: Cursor<Storage = B, Diff = Diff>,
C::Time: Copy,
C::ValOwn: 'static,
F: AsyncFnMut(
&<C::KeyContainer as BatchContainer>::Owned,
C::Time,
DiffPair<C::ValOwn>,
) -> Result<(), E>,
{
for (key, timed_pairs) in iter_diff_pairs(batch, lower, upper) {
for (time, pair) in timed_pairs {
on_diff_pair(&key, time, pair).await?;
}
}

Ok(())
}

/// Walks `batch` and emits, for each key, the `DiffPair`s at each timestamp.
///
/// Ignores updates outside the specified time range. Inclusive 'lower', exclusive 'upper'.
///
/// Within a key, diffs are partitioned by sign into retractions (befores) and
/// insertions (afters), sorted by timestamp, and zipped into `DiffPair`s via a
/// merge-join. Pairs are emitted in ascending timestamp order for a given key;
/// no ordering is guaranteed across keys. Callers are responsible for tracking
/// `(key, timestamp)` boundaries themselves if they need to detect groups
/// with more than one pair (e.g., for primary-key violation checks).
pub fn for_each_diff_pair<B, C, F>(batch: &B, mut on_diff_pair: F)
pub fn iter_diff_pairs<B, C>(
batch: &B,
lower: Option<Antichain<C::Time>>,
upper: Option<Antichain<C::Time>>,
) -> impl Iterator<
Item = (
<C::KeyContainer as BatchContainer>::Owned,
Vec<(C::Time, DiffPair<C::ValOwn>)>,
),
>
where
B: BatchReader<Time = C::Time> + Navigable<Cursor = C>,
C: Cursor<Storage = B, Diff = Diff>,
C::Time: Copy,
C::ValOwn: 'static,
F: FnMut(&<C::KeyContainer as BatchContainer>::Owned, C::Time, DiffPair<C::ValOwn>),
{
let mut befores: Vec<(C::Time, C::ValOwn, usize)> = vec![];
let mut afters: Vec<(C::Time, C::ValOwn, usize)> = vec![];

let mut cursor = batch.cursor();
while cursor.key_valid(batch) {
let k = cursor.key(batch);

// Partition updates at this key into retractions (befores) and
// insertions (afters).
while cursor.val_valid(batch) {
let v = cursor.val(batch);
cursor.map_times(batch, |t, diff| {
let diff = C::owned_diff(diff);
let update = (
C::owned_time(t),
C::owned_val(v),
usize::cast_from(diff.unsigned_abs()),
);
if diff < Diff::ZERO {
befores.push(update);
} else {
afters.push(update);
}
});
cursor.step_val(batch);
}
iter::from_fn(move || {
while cursor.key_valid(batch) {
let k = cursor.key(batch);

befores.sort_by_key(|(t, _v, _diff)| *t);
afters.sort_by_key(|(t, _v, _diff)| *t);

// The use of `repeat_n()` here is a bit subtle, but load bearing.
// In the common case, cnt = 1, and we want to avoid cloning the value if possible. In the naive
// implementation, we might use `iter::repeat((t, v)).take(cnt)`, but that would clone `v` `cnt` times even
// when `cnt = 1`. By contrast, `repeat_n((t, v), cnt)` will return the original `(t, v)` when `cnt = 1`,
// and only clone when `cnt > 1`.
let fan_out = |(t, v, cnt): (C::Time, C::ValOwn, usize)| iter::repeat_n((t, v), cnt);
let befores_iter = befores.drain(..).flat_map(fan_out);
let afters_iter = afters.drain(..).flat_map(fan_out);

let key_owned = <C::KeyContainer as BatchContainer>::into_owned(k);

for pair in
itertools::merge_join_by(befores_iter, afters_iter, |(t1, _v1), (t2, _v2)| t1.cmp(t2))
{
let (t, before, after) = match pair {
EitherOrBoth::Both((t, before), (_t, after)) => (t, Some(before), Some(after)),
EitherOrBoth::Left((t, before)) => (t, Some(before), None),
EitherOrBoth::Right((t, after)) => (t, None, Some(after)),
};
on_diff_pair(&key_owned, t, DiffPair { before, after });
}
// Partition updates at this key into retractions (befores) and
// insertions (afters).
while cursor.val_valid(batch) {
let v = cursor.val(batch);
cursor.map_times(batch, |t, diff| {
let t = C::owned_time(t);
if lower.as_ref().is_some_and(|lower| !lower.less_equal(&t))
|| upper.as_ref().is_some_and(|upper| upper.less_equal(&t))
{
// This record is outside the specified time interval. Ignore it.
return;
}
let diff = C::owned_diff(diff);
let update = (t, C::owned_val(v), usize::cast_from(diff.unsigned_abs()));
if diff < Diff::ZERO {
befores.push(update);
} else {
afters.push(update);
}
});
cursor.step_val(batch);
}

cursor.step_key(batch);
}
if befores.is_empty() && afters.is_empty() {
// No records at this key. Try the next one.
cursor.step_key(batch);
continue;
}

befores.sort_by_key(|(t, _v, _diff)| *t);
afters.sort_by_key(|(t, _v, _diff)| *t);

// The use of `repeat_n()` here is intentional.
// Typically, cnt = 1, and `repeat_n((t, v), cnt)` will return the original `(t, v)`.
// `iter::repeat((t, v)).take(cnt)` would clone `v` `cnt` times even when `cnt = 1`.
let fan_out = |(t, v, cnt): (C::Time, C::ValOwn, usize)| iter::repeat_n((t, v), cnt);
let befores_iter = befores.drain(..).flat_map(fan_out);
let afters_iter = afters.drain(..).flat_map(fan_out);

let key_owned = <C::KeyContainer as BatchContainer>::into_owned(k);

let pairs =
itertools::merge_join_by(befores_iter, afters_iter, |(t1, _v1), (t2, _v2)| {
t1.cmp(t2)
})
.map(|pair| {
let (t, before, after) = match pair {
EitherOrBoth::Both((t, before), (_t, after)) => {
(t, Some(before), Some(after))
}
EitherOrBoth::Left((t, before)) => (t, Some(before), None),
EitherOrBoth::Right((t, after)) => (t, None, Some(after)),
};
(t, DiffPair { before, after })
})
.collect();

cursor.step_key(batch);
return Some((key_owned, pairs));
}
None
})
}

// NOTE(benesch): statically allocating transient IDs for the
Expand Down Expand Up @@ -192,6 +266,36 @@ mod tests {
out
}

/// Collects `iter_diff_pairs` output into a per-key list, sorted by key
/// for deterministic assertion. Pair order within a key is preserved.
fn collect_bounded_diff_pairs<B, C>(
batch: &B,
lower: Option<u64>,
upper: Option<u64>,
) -> Vec<(String, Vec<(u64, Option<String>, Option<String>)>)>
where
B: BatchReader<Time = u64> + Navigable<Cursor = C>,
C: Cursor<Storage = B, Diff = Diff, Time = u64>,
C::ValOwn: 'static + Into<String>,
<C::KeyContainer as BatchContainer>::Owned: Into<String>,
{
let mut out: Vec<_> = iter_diff_pairs(
batch,
lower.map(Antichain::from_elem),
upper.map(Antichain::from_elem),
)
.map(|(k, pairs)| {
let pairs = pairs
.into_iter()
.map(|(t, dp)| (t, dp.before.map(Into::into), dp.after.map(Into::into)))
.collect();
(k.into(), pairs)
})
.collect();
out.sort();
out
}

#[mz_ore::test]
fn single_insertion() {
let batch = batch_from_tuples(vec![(("k1".into(), "v1".into()), 5, Diff::ONE)], 6);
Expand Down Expand Up @@ -299,4 +403,111 @@ mod tests {
]
);
}

#[mz_ore::test]
fn pairs_are_grouped_by_key() {
// One item per key, pairs within a key in ascending timestamp order.
let batch = batch_from_tuples(
vec![
(("k1".into(), "v1".into()), 5, Diff::ONE),
(("k1".into(), "v1".into()), 10, -Diff::ONE),
(("k1".into(), "v2".into()), 10, Diff::ONE),
(("k2".into(), "v3".into()), 7, Diff::ONE),
],
11,
);
let items = collect_bounded_diff_pairs(&batch, None, None);
assert_eq!(
items,
vec![
(
"k1".into(),
vec![
(5, None, Some("v1".into())),
(10, Some("v1".into()), Some("v2".into())),
]
),
("k2".into(), vec![(7, None, Some("v3".into()))]),
]
);
}

#[mz_ore::test]
fn lower_bound_is_inclusive() {
let batch = batch_from_tuples(
vec![
(("k1".into(), "v1".into()), 4, Diff::ONE),
(("k1".into(), "v2".into()), 5, Diff::ONE),
],
6,
);
let items = collect_bounded_diff_pairs(&batch, Some(5), None);
assert_eq!(
items,
vec![("k1".into(), vec![(5, None, Some("v2".into()))])]
);
}

#[mz_ore::test]
fn upper_bound_is_exclusive() {
let batch = batch_from_tuples(
vec![
(("k1".into(), "v1".into()), 5, Diff::ONE),
(("k1".into(), "v2".into()), 6, Diff::ONE),
],
7,
);
let items = collect_bounded_diff_pairs(&batch, None, Some(6));
assert_eq!(
items,
vec![("k1".into(), vec![(5, None, Some("v1".into()))])]
);
}

#[mz_ore::test]
fn bounds_filter_within_a_key() {
// The window [6, 11) drops the t=5 insertion but keeps the t=10
// retraction/insertion, which still pair with each other.
let batch = batch_from_tuples(
vec![
(("k1".into(), "v1".into()), 5, Diff::ONE),
(("k1".into(), "v1".into()), 10, -Diff::ONE),
(("k1".into(), "v2".into()), 10, Diff::ONE),
],
11,
);
let items = collect_bounded_diff_pairs(&batch, Some(6), Some(11));
assert_eq!(
items,
vec![(
"k1".into(),
vec![(10, Some("v1".into()), Some("v2".into()))]
)]
);
}

#[mz_ore::test]
fn fully_filtered_key_emits_no_item() {
// All of k1's updates fall outside the bounds, so k1 must not appear
// at all, not even as a key with an empty pair list.
let batch = batch_from_tuples(
vec![
(("k1".into(), "v1".into()), 3, Diff::ONE),
(("k2".into(), "v2".into()), 5, Diff::ONE),
],
6,
);
let items = collect_bounded_diff_pairs(&batch, Some(5), None);
assert_eq!(
items,
vec![("k2".into(), vec![(5, None, Some("v2".into()))])]
);
}

#[mz_ore::test]
fn bounds_filter_everything() {
let batch = batch_from_tuples(vec![(("k1".into(), "v1".into()), 5, Diff::ONE)], 6);
let items = collect_bounded_diff_pairs(&batch, Some(6), None);
assert_eq!(items, vec![]);
}
}
Loading
Loading