1use std::{
4 collections::{BTreeMap, BTreeSet, HashMap},
5 iter,
6};
7
8use flatbuffers::VerifierOptions;
9use itertools::{Either, Itertools as _};
10
11use crate::{
12 ChunkIndices, IcechunkFormatErrorKind, IcechunkResult, Move, NodeId, Path,
13 SnapshotId,
14 flatbuffers::generated::{self, MoveOperation, MoveOperationArgs},
15};
16use icechunk_types::ICResultExt as _;
17
18#[derive(Clone, Debug, PartialEq, Default)]
19pub struct TransactionLog {
20 buffer: Vec<u8>,
21}
22
23impl TransactionLog {
24 #[expect(clippy::too_many_arguments)]
27 pub fn new_from_parts(
28 id: &SnapshotId,
29 sorted_new_groups: impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator,
30 sorted_new_arrays: impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator,
31 sorted_deleted_groups: impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator,
32 sorted_deleted_arrays: impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator,
33 sorted_updated_groups: impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator,
34 sorted_updated_arrays: impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator,
35 sorted_updated_chunks: impl ExactSizeIterator<
36 Item = (NodeId, impl Iterator<Item = ChunkIndices>),
37 > + DoubleEndedIterator,
38 sorted_moves: impl Iterator<Item = Move>,
39 ) -> Self {
40 let mut builder = flatbuffers::FlatBufferBuilder::with_capacity(1_024 * 1_024);
42
43 let new_groups = Some(builder.create_vector_from_iter(
45 sorted_new_groups.map(|id| generated::ObjectId8::new(&id.0)),
46 ));
47 let new_arrays = Some(builder.create_vector_from_iter(
48 sorted_new_arrays.map(|id| generated::ObjectId8::new(&id.0)),
49 ));
50 let deleted_groups = Some(builder.create_vector_from_iter(
51 sorted_deleted_groups.map(|id| generated::ObjectId8::new(&id.0)),
52 ));
53 let deleted_arrays = Some(builder.create_vector_from_iter(
54 sorted_deleted_arrays.map(|id| generated::ObjectId8::new(&id.0)),
55 ));
56 let updated_groups = Some(builder.create_vector_from_iter(
57 sorted_updated_groups.map(|id| generated::ObjectId8::new(&id.0)),
58 ));
59 let updated_arrays = Some(builder.create_vector_from_iter(
60 sorted_updated_arrays.map(|id| generated::ObjectId8::new(&id.0)),
61 ));
62
63 let id = generated::ObjectId12::new(&id.0);
64 let id = Some(&id);
65 let updated_chunks = sorted_updated_chunks
67 .map(|(node_id, chunks)| {
68 let node_id = generated::ObjectId8::new(&node_id.0);
69 let node_id = Some(&node_id);
70 let chunks = chunks
71 .map(|indices| {
72 let coords = Some(builder.create_vector(indices.0.as_slice()));
73 generated::ChunkIndices::create(
74 &mut builder,
75 &generated::ChunkIndicesArgs { coords },
76 )
77 })
78 .collect::<Vec<_>>();
79 let chunks = Some(builder.create_vector(chunks.as_slice()));
80 generated::ArrayUpdatedChunks::create(
81 &mut builder,
82 &generated::ArrayUpdatedChunksArgs { node_id, chunks },
83 )
84 })
85 .collect::<Vec<_>>();
86 let updated_chunks = builder.create_vector(updated_chunks.as_slice());
87 let updated_chunks = Some(updated_chunks);
88
89 let moved_nodes: Vec<_> = sorted_moves
90 .map(|Move { from, to, node_id, node_type }| {
91 let from = builder.create_string(from.to_string().as_str());
92 let to = builder.create_string(to.to_string().as_str());
93 let node_id = generated::ObjectId8::new(&node_id.0);
94 let node_id = Some(&node_id);
95 let node_type: generated::NodeType = node_type.into();
96 let args = MoveOperationArgs {
97 from: Some(from),
98 to: Some(to),
99 node_id,
100 node_type,
101 };
102 MoveOperation::create(&mut builder, &args)
103 })
104 .collect();
105
106 let moved_nodes = Some(builder.create_vector(moved_nodes.as_slice()));
107
108 let tx = generated::TransactionLog::create(
109 &mut builder,
110 &generated::TransactionLogArgs {
111 id,
112 new_groups,
113 new_arrays,
114 deleted_groups,
115 deleted_arrays,
116 updated_groups,
117 updated_arrays,
118 updated_chunks,
119 moved_nodes,
120 ..Default::default()
121 },
122 );
123
124 builder.finish(tx, Some("Ichk"));
125 let (mut buffer, offset) = builder.collapse();
126 buffer.drain(0..offset);
127 buffer.shrink_to_fit();
128 Self { buffer }
129 }
130
131 pub fn verify_buffer(buffer: &[u8]) -> IcechunkResult<()> {
133 let _ = flatbuffers::root_with_opts::<generated::TransactionLog<'_>>(
134 &ROOT_OPTIONS,
135 buffer,
136 )
137 .capture()?;
138 Ok(())
139 }
140
141 pub fn from_buffer(buffer: Vec<u8>) -> IcechunkResult<Self> {
142 Self::verify_buffer(buffer.as_slice())?;
143 Ok(Self { buffer })
144 }
145
146 pub fn new_groups(
147 &self,
148 ) -> impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator + '_ {
149 self.root().new_groups().iter().map(From::from)
150 }
151
152 pub fn new_arrays(
153 &self,
154 ) -> impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator + '_ {
155 self.root().new_arrays().iter().map(From::from)
156 }
157
158 pub fn deleted_groups(
159 &self,
160 ) -> impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator + '_ {
161 self.root().deleted_groups().iter().map(From::from)
162 }
163
164 pub fn deleted_arrays(
165 &self,
166 ) -> impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator + '_ {
167 self.root().deleted_arrays().iter().map(From::from)
168 }
169
170 pub fn updated_groups(
171 &self,
172 ) -> impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator + '_ {
173 self.root().updated_groups().iter().map(From::from)
174 }
175
176 pub fn updated_arrays(
177 &self,
178 ) -> impl ExactSizeIterator<Item = NodeId> + DoubleEndedIterator + '_ {
179 self.root().updated_arrays().iter().map(From::from)
180 }
181
182 pub fn updated_chunks(
183 &self,
184 ) -> impl Iterator<Item = (NodeId, impl Iterator<Item = ChunkIndices> + '_)> + '_
185 {
186 self.root().updated_chunks().iter().map(|arr_chunks| {
187 let id: NodeId = arr_chunks.node_id().into();
188 let chunks = arr_chunks.chunks().iter().map(|idx| idx.into());
189 (id, chunks)
190 })
191 }
192
193 pub fn moves(&self) -> impl Iterator<Item = IcechunkResult<Move>> + '_ {
194 let it = match self.root().moved_nodes() {
195 Some(it) => Either::Left(it.iter()),
196 None => Either::Right(iter::empty()),
197 };
198 it.map(|m| {
199 let Some(from) = m.from() else {
200 return Err(IcechunkFormatErrorKind::MissingRequiredField("from".into()))
201 .capture();
202 };
203 let from = match Path::new(from) {
204 Ok(from) => from,
205 Err(e) => return Err(IcechunkFormatErrorKind::Path(e)).capture(),
206 };
207
208 let Some(to) = m.to() else {
209 return Err(IcechunkFormatErrorKind::MissingRequiredField("to".into()))
210 .capture();
211 };
212 let to = match Path::new(to) {
213 Ok(to) => to,
214 Err(e) => return Err(IcechunkFormatErrorKind::Path(e)).capture(),
215 };
216
217 let Some(node_id) = m.node_id() else {
218 return Err(IcechunkFormatErrorKind::MissingRequiredField(
219 "node_id".into(),
220 ))
221 .capture();
222 };
223
224 let node_type = m.node_type().try_into().capture()?;
225 Ok(Move { from, to, node_id: node_id.into(), node_type })
226 })
227 }
228
229 pub fn updated_chunks_for(
230 &self,
231 node: &NodeId,
232 ) -> impl Iterator<Item = ChunkIndices> + '_ + use<'_> {
233 let arr = self
234 .root()
235 .updated_chunks()
236 .lookup_by_key(node.0, |a, b| a.node_id().0.cmp(b));
237
238 match arr {
239 Some(arr) => Either::Left(arr.chunks().iter().map(From::from)),
240 None => Either::Right(iter::empty()),
241 }
242 }
243
244 pub fn updated_chunks_counts(
245 &self,
246 ) -> impl Iterator<Item = (NodeId, u64)> + '_ + use<'_> {
247 self.root().updated_chunks().iter().map(|arr_chunks| {
248 let id: NodeId = arr_chunks.node_id().into();
249 let n = arr_chunks.chunks().len();
250 (id, n as u64)
251 })
252 }
253
254 pub fn group_created(&self, id: &NodeId) -> bool {
255 self.root().new_groups().lookup_by_key(id.0, |a, b| a.0.cmp(b)).is_some()
256 }
257
258 pub fn array_created(&self, id: &NodeId) -> bool {
259 self.root().new_arrays().lookup_by_key(id.0, |a, b| a.0.cmp(b)).is_some()
260 }
261
262 pub fn group_deleted(&self, id: &NodeId) -> bool {
263 self.root().deleted_groups().lookup_by_key(id.0, |a, b| a.0.cmp(b)).is_some()
264 }
265
266 pub fn array_deleted(&self, id: &NodeId) -> bool {
267 self.root().deleted_arrays().lookup_by_key(id.0, |a, b| a.0.cmp(b)).is_some()
268 }
269
270 pub fn group_updated(&self, id: &NodeId) -> bool {
271 self.root().updated_groups().lookup_by_key(id.0, |a, b| a.0.cmp(b)).is_some()
272 }
273
274 pub fn array_updated(&self, id: &NodeId) -> bool {
275 self.root().updated_arrays().lookup_by_key(id.0, |a, b| a.0.cmp(b)).is_some()
276 }
277
278 pub fn chunks_updated(&self, id: &NodeId) -> bool {
279 self.root()
280 .updated_chunks()
281 .lookup_by_key(id.0, |a, b| a.node_id().0.cmp(b))
282 .is_some()
283 }
284
285 #[expect(unsafe_code)]
286 fn root(&self) -> generated::TransactionLog<'_> {
287 unsafe {
291 flatbuffers::root_unchecked::<generated::TransactionLog<'_>>(&self.buffer)
292 }
293 }
294
295 pub fn bytes(&self) -> &[u8] {
296 self.buffer.as_slice()
297 }
298
299 pub fn len(&self) -> usize {
300 let root = self.root();
301 root.new_groups().len()
302 + root.new_arrays().len()
303 + root.deleted_groups().len()
304 + root.deleted_arrays().len()
305 + root.updated_groups().len()
306 + root.updated_arrays().len()
307 + root.updated_chunks().iter().map(|s| s.chunks().len()).sum::<usize>()
308 + root.moved_nodes().map(|v| v.len()).unwrap_or_default()
309 }
310
311 #[must_use]
312 pub fn is_empty(&self) -> bool {
313 self.len() == 0
314 }
315
316 pub fn has_moves(&self) -> bool {
317 self.root().moved_nodes().map(|v| !v.is_empty()).unwrap_or(false)
318 }
319
320 pub fn merge<'a, T: IntoIterator<Item = &'a TransactionLog>>(
321 id: &SnapshotId,
322 iter: T,
323 ) -> IcechunkResult<Self> {
324 let txs = Vec::from_iter(iter);
325
326 let mut builder = flatbuffers::FlatBufferBuilder::with_capacity(1_024 * 1_024);
328
329 let new_groups = {
330 let new_groups =
331 BTreeSet::from_iter(txs.iter().flat_map(|tx| tx.new_groups()));
332 let new_groups = Vec::from_iter(
333 new_groups.into_iter().map(|id| generated::ObjectId8::new(&id.0)),
334 );
335 Some(builder.create_vector(new_groups.as_slice()))
336 };
337 let new_arrays = {
338 let new_arrays =
339 BTreeSet::from_iter(txs.iter().flat_map(|tx| tx.new_arrays()));
340 let new_arrays = Vec::from_iter(
341 new_arrays.into_iter().map(|id| generated::ObjectId8::new(&id.0)),
342 );
343 Some(builder.create_vector(new_arrays.as_slice()))
344 };
345 let deleted_groups = {
346 let deleted_groups =
347 BTreeSet::from_iter(txs.iter().flat_map(|tx| tx.deleted_groups()));
348 let deleted_groups = Vec::from_iter(
349 deleted_groups.into_iter().map(|id| generated::ObjectId8::new(&id.0)),
350 );
351 Some(builder.create_vector(deleted_groups.as_slice()))
352 };
353 let deleted_arrays = {
354 let deleted_arrays =
355 BTreeSet::from_iter(txs.iter().flat_map(|tx| tx.deleted_arrays()));
356 let deleted_arrays = Vec::from_iter(
357 deleted_arrays.into_iter().map(|id| generated::ObjectId8::new(&id.0)),
358 );
359 Some(builder.create_vector(deleted_arrays.as_slice()))
360 };
361 let updated_groups = {
362 let updated_groups =
363 BTreeSet::from_iter(txs.iter().flat_map(|tx| tx.updated_groups()));
364 let updated_groups = Vec::from_iter(
365 updated_groups.into_iter().map(|id| generated::ObjectId8::new(&id.0)),
366 );
367 Some(builder.create_vector(updated_groups.as_slice()))
368 };
369 let updated_arrays = {
370 let updated_arrays =
371 BTreeSet::from_iter(txs.iter().flat_map(|tx| tx.updated_arrays()));
372 let updated_arrays = Vec::from_iter(
373 updated_arrays.into_iter().map(|id| generated::ObjectId8::new(&id.0)),
374 );
375 Some(builder.create_vector(updated_arrays.as_slice()))
376 };
377 let updated_chunks = {
378 let updated_chunks = txs.iter().fold(BTreeMap::new(), |res, tx| {
379 tx.updated_chunks().fold(res, |mut res, (node_id, chunks_it)| {
380 let set: &mut BTreeSet<_> = res.entry(node_id).or_default();
381 set.extend(chunks_it);
382 res
383 })
384 });
385
386 let updated_chunks = updated_chunks
387 .into_iter()
388 .map(|(node_id, chunks)| {
389 let node_id = generated::ObjectId8::new(&node_id.0);
390 let node_id = Some(&node_id);
391 let chunks = chunks
392 .into_iter()
393 .map(|indices| {
394 let coords =
395 Some(builder.create_vector(indices.0.as_slice()));
396 generated::ChunkIndices::create(
397 &mut builder,
398 &generated::ChunkIndicesArgs { coords },
399 )
400 })
401 .collect::<Vec<_>>();
402 let chunks = Some(builder.create_vector(chunks.as_slice()));
403 generated::ArrayUpdatedChunks::create(
404 &mut builder,
405 &generated::ArrayUpdatedChunksArgs { node_id, chunks },
406 )
407 })
408 .collect::<Vec<_>>();
409
410 let updated_chunks = builder.create_vector(updated_chunks.as_slice());
411 Some(updated_chunks)
412 };
413
414 let id = generated::ObjectId12::new(&id.0);
415 let id = Some(&id);
416
417 let mut moved_map: HashMap<NodeId, Move> = Default::default();
440 let moved_nodes = {
441 for mv in txs.iter().flat_map(|tx| tx.moves()) {
442 let Move { from, to, node_id, node_type } = mv?;
443 moved_map
448 .entry(node_id.clone())
449 .and_modify(|m| m.to = to.clone())
450 .or_insert_with(|| Move { to, from, node_id, node_type });
451 }
452
453 moved_map.retain(|_, mv| mv.from != mv.to);
456
457 debug_assert!({
460 let node_id_len = moved_map.len();
461
462 let (from, to): (BTreeSet<&Path>, BTreeSet<&Path>) = moved_map
463 .values()
464 .map(|Move { to, from, .. }| (from, to))
465 .multiunzip();
466
467 (from.len() == to.len()) && (to.len() == node_id_len)
468 });
469
470 let moved_nodes: Vec<_> = moved_map
471 .into_values()
472 .sorted_by(|a, b| a.to.cmp(&b.to))
475 .map(|Move { to, from, node_id, node_type }| {
476 let from = builder.create_string(from.to_string().as_str());
477 let to = builder.create_string(to.to_string().as_str());
478 let node_id = generated::ObjectId8::new(&node_id.0);
479 let node_id = Some(&node_id);
480 let node_type: generated::NodeType = node_type.into();
481 let args = MoveOperationArgs {
482 from: Some(from),
483 to: Some(to),
484 node_id,
485 node_type,
486 };
487 MoveOperation::create(&mut builder, &args)
488 })
489 .collect();
490 Some(builder.create_vector(moved_nodes.as_slice()))
491 };
492
493 let tx = generated::TransactionLog::create(
494 &mut builder,
495 &generated::TransactionLogArgs {
496 id,
497 new_groups,
498 new_arrays,
499 deleted_groups,
500 deleted_arrays,
501 updated_groups,
502 updated_arrays,
503 updated_chunks,
504 moved_nodes,
505 ..Default::default()
506 },
507 );
508
509 builder.finish(tx, Some("Ichk"));
510 let (mut buffer, offset) = builder.collapse();
511 buffer.drain(0..offset);
512 buffer.shrink_to_fit();
513 Ok(Self { buffer })
514 }
515}
516
517static ROOT_OPTIONS: VerifierOptions = VerifierOptions {
518 max_depth: 64,
519 max_tables: 50_000_000,
520 max_apparent_size: crate::serializers::MAX_APPARENT_SIZE_INFLATION
526 .saturating_mul(flatbuffers::FLATBUFFERS_MAX_BUFFER_SIZE),
527 ignore_missing_null_terminator: true,
528};
529
530#[cfg(test)]
532pub(crate) fn root_options() -> &'static VerifierOptions {
533 &ROOT_OPTIONS
534}
535
536#[cfg(any())]
539mod tests {
540 use std::collections::HashSet;
541
542 use bytes::Bytes;
543 use itertools::Itertools as _;
544
545 use crate::{
546 change_set::{ArrayData, ChangeSet, transaction_log_from_change_set},
547 format::{
548 ChunkIndices, NodeId, SnapshotId, manifest::ChunkPayload,
549 snapshot::ArrayShape, transaction_log::TransactionLog,
550 },
551 };
552
553 #[icechunk_macros::test]
554 fn test_merge() -> Result<(), Box<dyn std::error::Error>> {
555 let mut cs1 = ChangeSet::for_edits();
556 let added_group = NodeId::random();
557 let added_array = NodeId::random();
558 let deleted_group = NodeId::random();
559 let deleted_array = NodeId::random();
560 let updated_group = NodeId::random();
561 let chunk_added = NodeId::random();
562 cs1.add_group("/g1".try_into().unwrap(), added_group.clone(), Bytes::new())?;
563 cs1.delete_group("/g2".try_into().unwrap(), &deleted_group)?;
564 cs1.add_array(
565 "/a1".try_into().unwrap(),
566 added_array.clone(),
567 ArrayData {
568 shape: ArrayShape::new([(0, 10)]).unwrap(),
569 dimension_names: None,
570 user_data: Bytes::new(),
571 },
572 )?;
573 cs1.delete_array("/a2".try_into().unwrap(), &deleted_array)?;
574 cs1.update_group(&updated_group, &"/g3".try_into().unwrap(), Bytes::new())?;
575 cs1.set_chunk_ref(
576 chunk_added.clone(),
577 ChunkIndices(vec![0]),
578 Some(ChunkPayload::Inline(Bytes::new())),
579 )?;
580
581 let t1 = transaction_log_from_change_set(&SnapshotId::random(), &cs1);
582 let t2 = transaction_log_from_change_set(&SnapshotId::random(), &cs1);
583
584 let tx = TransactionLog::merge(&SnapshotId::random(), [&t1, &t2]);
585 assert!(tx.new_groups().eq([added_group.clone()]));
586 assert!(tx.new_arrays().eq([added_array.clone()]));
587 assert!(tx.deleted_groups().eq([deleted_group.clone()]));
588 assert!(tx.deleted_arrays().eq([deleted_array.clone()]));
589 assert!(tx.updated_groups().eq([updated_group.clone()]));
590 let chunks =
591 Vec::from_iter(tx.updated_chunks().map(|(id, it)| (id, Vec::from_iter(it))));
592 assert_eq!(chunks, vec![(chunk_added.clone(), vec![ChunkIndices(vec![0])])]);
593 assert_eq!(
594 tx.updated_chunks_counts().collect::<Vec<_>>(),
595 vec![(chunk_added.clone(), 1)]
596 );
597
598 let added_group2 = NodeId::random();
599 let deleted_group2 = NodeId::random();
600 let deleted_array2 = NodeId::random();
601 let updated_group2 = NodeId::random();
602 let chunk_added2 = NodeId::random();
603 let mut cs2 = ChangeSet::for_edits();
604 cs2.add_group("/g1".try_into().unwrap(), added_group2.clone(), Bytes::new())?;
605 cs2.delete_group("/g2".try_into().unwrap(), &deleted_group2)?;
606 cs2.add_array(
607 "/a1".try_into().unwrap(),
608 added_array.clone(),
609 ArrayData {
610 shape: ArrayShape::new([(0, 10)]).unwrap(),
611 dimension_names: None,
612 user_data: Bytes::new(),
613 },
614 )?;
615 cs2.delete_array("/a2".try_into().unwrap(), &deleted_array2)?;
616 cs2.update_group(&updated_group2, &"/g3".try_into().unwrap(), Bytes::new())?;
617 cs2.set_chunk_ref(chunk_added.clone(), ChunkIndices(vec![0]), None)?;
618 cs2.set_chunk_ref(
619 chunk_added.clone(),
620 ChunkIndices(vec![1]),
621 Some(ChunkPayload::Inline(Bytes::new())),
622 )?;
623 cs2.set_chunk_ref(chunk_added.clone(), ChunkIndices(vec![42]), None)?;
624 cs2.set_chunk_ref(
625 chunk_added2.clone(),
626 ChunkIndices(vec![7]),
627 Some(ChunkPayload::Inline(Bytes::new())),
628 )?;
629
630 let t3 = transaction_log_from_change_set(&SnapshotId::random(), &cs2);
631 let tx_id = SnapshotId::random();
632 let tx = TransactionLog::merge(&tx_id, [&t1, &t2, &t3]);
633
634 assert!(
635 tx.new_groups()
636 .sorted()
637 .eq([added_group.clone(), added_group2.clone()].into_iter().sorted())
638 );
639 assert!(tx.new_arrays().eq([added_array.clone()]));
640 assert!(
641 tx.deleted_groups().sorted().eq([
642 deleted_group.clone(),
643 deleted_group2.clone()
644 ]
645 .into_iter()
646 .sorted())
647 );
648 assert!(
649 tx.deleted_arrays().sorted().eq([
650 deleted_array.clone(),
651 deleted_array2.clone()
652 ]
653 .into_iter()
654 .sorted())
655 );
656 assert!(
657 tx.updated_groups().sorted().eq([
658 updated_group.clone(),
659 updated_group2.clone()
660 ]
661 .into_iter()
662 .sorted())
663 );
664 let chunks = HashSet::from_iter(
665 tx.updated_chunks().map(|(id, it)| (id, Vec::from_iter(it))),
666 );
667 assert_eq!(
668 chunks,
669 HashSet::from([
670 (
671 chunk_added.clone(),
672 vec![
673 ChunkIndices(vec![0]),
674 ChunkIndices(vec![1]),
675 ChunkIndices(vec![42])
676 ]
677 ),
678 (chunk_added2.clone(), vec![ChunkIndices(vec![7]),])
679 ])
680 );
681
682 assert_eq!(
683 tx.updated_chunks_counts().collect::<HashSet<_>>(),
684 HashSet::from([(chunk_added.clone(), 3), (chunk_added2, 1)])
685 );
686 Ok(())
687 }
688}