simple refactor StreamLimit
@@ -36,60 +36,42 @@ impl LimitStream {
pub fn limit(&mut self, block: &DataBlock) -> Result<Option<DataBlock>> {
let rows = block.num_rows();
+ let limit = self.limit.unwrap_or(usize::MAX - self.offset);
+ let offset = self.offset;
+ let current = self.current;
+
// There are two intervals:
// r1: [current, current + rows)
// r2: [offset, offset + limit), note that limit may be infinite.
//
// There are 6 possible relationships between r1 and r2:
- //
- // TODO: We can clean some codes while RFC 2497-if-let-chains stable.
- // https://github.com/rust-lang/rust/issues/53668
- if self.current + rows <= self.offset {
+ if current + rows <= offset {
// case 1: no overlap. r1.r <= r2.l.
// output nothing.
self.current += rows;
- Ok(None)
- } else if self
- .limit
- .map(|limit| self.offset + limit == self.current)
- .unwrap_or(false)
- {
+ Ok(Some(DataBlock::empty_with_schema(block.schema().clone())))
+ } else if offset + limit == current {
// case 2: no overlap. r2.r == r1.l.
// output nothing.
Ok(None)
- } else if self.current <= self.offset
- && self
- .limit
- .map(|limit| self.current + rows >= self.offset + limit)
- .unwrap_or(false)
- {
+ } else if current <= offset && current + rows >= offset + limit {
// case 3: r1 contains r2.
// output r2.
- self.current = self.offset + self.limit.unwrap();
- Ok(Some(block.slice(self.offset, self.limit.unwrap())))
- } else if self.offset <= self.current
- && self
- .limit
- .map(|limit| self.offset + limit >= self.current + rows)
- .unwrap_or(true)
- {
+ self.current += offset + limit;
+ Ok(Some(block.slice(offset, limit)))
+ } else if offset <= current && offset + limit >= current + rows {
// case 4: r2 contains r1.
// output r1.
self.current += rows;
Ok(Some(block.clone()))
- } else if self.current <= self.offset {
+ } else if current <= offset {
// case 5: overlap and r1 in the left. r1.l <= r2.l
self.current += rows;
- Ok(Some(
- block.slice(self.offset, self.current + rows - self.offset),
- ))
+ Ok(Some(block.slice(offset, current + rows - offset)))
} else {
// case 6: overlap and r2 in the left.
self.current = self.offset + self.limit.unwrap();
- Ok(Some(block.slice(
- self.current,
- self.offset + self.limit.unwrap() - self.current,
- )))
+ Ok(Some(block.slice(current, offset + limit - current)))
}
}
}
