From 9fa438206139391600a38f1a9259d63b3f5c0f40 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 26 Aug 2024 15:06:22 -0700 Subject: [PATCH] [BUG] Fix limit on MicroPartition with multiple Tables --- src/daft-micropartition/src/ops/slice.rs | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/daft-micropartition/src/ops/slice.rs b/src/daft-micropartition/src/ops/slice.rs index fa6cb858f1..01de180063 100644 --- a/src/daft-micropartition/src/ops/slice.rs +++ b/src/daft-micropartition/src/ops/slice.rs @@ -27,17 +27,18 @@ impl MicroPartition { } let tab_rows = tab.len(); - if offset_so_far > 0 && offset_so_far >= tab_rows { - offset_so_far -= tab_rows; - continue; - } - if offset_so_far == 0 && rows_needed >= tab_rows { + if offset_so_far >= tab_rows { + // we can skip the entire table + offset_so_far -= tab_rows; + } else if offset_so_far == 0 && rows_needed >= tab_rows { + // we can take the entire table slices_tables.push(tab.clone()); rows_needed -= tab_rows; } else { - let new_end = (rows_needed + offset_so_far).min(tab_rows); - let sliced = tab.slice(offset_so_far, new_end)?; + // we need to slice the table + let tab_end = (offset_so_far + rows_needed).min(tab_rows); + let sliced = tab.slice(offset_so_far, tab_end)?; offset_so_far = 0; rows_needed -= sliced.len(); slices_tables.push(sliced);