[query/datasources] move datablock generation to get_one_block
@@ -44,7 +44,7 @@ impl NumbersStream {
}
}
- fn try_get_one_block(&mut self) -> Result<Option<BlockRange>> {
+ fn try_get_one_block(&mut self) -> Result<Option<DataBlock>> {
if (self.block_index as usize) == self.blocks.len() {
let partitions = self.ctx.try_get_partitions(1)?;
if partitions.is_empty() {
@@ -85,23 +85,10 @@ impl NumbersStream {
let current = self.blocks[self.block_index].clone();
self.block_index += 1;
- Ok(Some(current))
- }
-}
-
-impl Stream for NumbersStream {
- type Item = Result<DataBlock>;
-
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- _: &mut Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- let current = self.try_get_one_block()?;
- Poll::Ready(match current {
- None => None,
- Some(current) if current.begin == current.end => None,
- Some(current) => {
+ Ok(if current.begin == current.end {
+ None
+ } else {
let v = (current.begin..current.end).collect::<Vec<u64>>();
let mut me = ManuallyDrop::new(v);
let byte_size = mem::size_of::<u64>();
@@ -123,9 +110,21 @@ impl Stream for NumbersStream {
self.schema.clone(),
vec![Arc::new(UInt64Array::from(arr_data))],
);
- Some(Ok(block))
- }
+ Some(block)
}
})
}
}
+
+impl Stream for NumbersStream {
+ type Item = Result<DataBlock>;
+
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ _: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let block = self.try_get_one_block()?;
+
+ Poll::Ready(block.map(Ok))
+ }
+}
