perf: split MaterializeIndex stream into batches

lancedb · Aug 21, 2024 · 35adc71 · 35adc71
1 parent a9c678d
commit 35adc71
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs
@@ -23,6 +23,7 @@ use lance_core::{
     },
     Error, Result, ROW_ID_FIELD,
 };
+use lance_datafusion::chunker::break_stream;
 use lance_index::{
     scalar::{
         expression::{ScalarIndexExpr, ScalarIndexLoader},
@@ -582,9 +583,14 @@ impl ExecutionPlan for MaterializeIndexExec {
             .then(|batch_fut| batch_fut.map_err(|err| err.into()))
             .boxed()
             as BoxStream<'static, datafusion::common::Result<RecordBatch>>;
-        Ok(Box::pin(RecordBatchStreamAdapter::new(
+        let stream = Box::pin(RecordBatchStreamAdapter::new(
             MATERIALIZE_INDEX_SCHEMA.clone(),
             stream,
+        ));
+        let stream = break_stream(stream, 50 * 1024);
+        Ok(Box::pin(RecordBatchStreamAdapter::new(
+            MATERIALIZE_INDEX_SCHEMA.clone(),
+            stream.map_err(|err| err.into()),
         )))
     }