use std::collections::{HashMap, hash_map::Entry};
use std::fmt::{Debug, Formatter};
use std::sync::Arc;
use arrow_array::builder::{ListBuilder, UInt32Builder};
use arrow_array::{BooleanArray, Float32Array, RecordBatch, UInt32Array, UInt64Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use datafusion::common::ScalarValue;
use datafusion::common::stats::Precision;
use datafusion::error::Result as DataFusionResult;
use datafusion::execution::TaskContext;
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
SendableRecordBatchStream, Statistics,
};
use datafusion_physical_expr::{EquivalenceProperties, PhysicalExprRef};
use futures::stream::{self, StreamExt};
use lance_core::{Error, Result};
use lance_index::scalar::inverted::DOC_INDEX_FIELD;
use super::super::builder::FtsQuery;
use super::{newest_pk_positions, scan_record_batch};
use crate::dataset::mem_wal::index::{
FtsEntry, FtsMemQuery, MemMatches, SearchContext, SearchOptions, search_cross_column,
};
use crate::dataset::mem_wal::memtable::scanner::exec::take_projected_columns;
use crate::dataset::mem_wal::scanner::exec::resolve_pk_indices;
use crate::dataset::mem_wal::write::{BatchStore, IndexStore};
pub const SCORE_COLUMN: &str = "_score";
#[derive(Debug, Clone)]
struct BatchRange {
start: usize,
end: usize,
batch_id: usize,
}
type FtsHit = (u64, Option<Vec<u32>>, f32);
type MaterializedFtsRows = (
Vec<Arc<dyn arrow_array::Array>>,
Vec<f32>,
Vec<u64>,
Vec<Option<Vec<u32>>>,
);
pub struct FtsIndexExec {
batch_store: Arc<BatchStore>,
indexes: Arc<IndexStore>,
query: FtsQuery,
readable_count: usize,
projection: Option<Vec<usize>>,
output_schema: SchemaRef,
properties: Arc<PlanProperties>,
metrics: ExecutionPlanMetricsSet,
batch_ranges: Vec<BatchRange>,
max_readable_row: Option<u64>,
with_row_id: bool,
with_doc_index: bool,
filter: Option<PhysicalExprRef>,
pk_columns: Option<Vec<String>>,
}
impl Debug for FtsIndexExec {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FtsIndexExec")
.field("columns", &self.query.columns())
.field("expr", &self.query.expr)
.field("readable_count", &self.readable_count)
.field("with_row_id", &self.with_row_id)
.finish()
}
}
impl FtsIndexExec {
pub fn new(
batch_store: Arc<BatchStore>,
indexes: Arc<IndexStore>,
query: FtsQuery,
readable_count: usize,
projection: Option<Vec<usize>>,
base_schema: SchemaRef,
with_row_id: bool,
) -> Result<Self> {
if !query.is_answered_by(&indexes) {
return Err(Error::invalid_input(format!(
"no full-text index answers this search over {:?}",
query.columns()
)));
}
let with_doc_index = query.document_granularity.is_list_element();
let mut fields: Vec<Field> = base_schema
.fields()
.iter()
.map(|f| f.as_ref().clone())
.collect();
if with_doc_index {
fields.push(DOC_INDEX_FIELD.clone());
}
fields.push(Field::new(SCORE_COLUMN, DataType::Float32, true));
if with_row_id {
fields.push(Field::new(lance_core::ROW_ID, DataType::UInt64, true));
}
let output_schema = Arc::new(Schema::new(fields));
let properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(output_schema.clone()),
Partitioning::UnknownPartitioning(1),
EmissionType::Incremental,
Boundedness::Bounded,
));
let mut batch_ranges = Vec::new();
let mut current_row = 0usize;
let mut max_readable_row_exclusive: u64 = 0;
for (batch_id, stored_batch) in batch_store.iter().enumerate() {
let batch_start = current_row;
let batch_end = current_row + stored_batch.num_rows;
batch_ranges.push(BatchRange {
start: batch_start,
end: batch_end,
batch_id,
});
if batch_id < readable_count {
max_readable_row_exclusive = batch_end as u64;
}
current_row = batch_end;
}
let max_readable_row = if max_readable_row_exclusive > 0 {
Some(max_readable_row_exclusive - 1)
} else {
None
};
Ok(Self {
batch_store,
indexes,
query,
readable_count,
projection,
output_schema,
properties,
metrics: ExecutionPlanMetricsSet::new(),
batch_ranges,
max_readable_row,
with_row_id,
with_doc_index,
filter: None,
pk_columns: None,
})
}
pub fn with_filter(mut self, filter: Option<PhysicalExprRef>) -> Self {
self.filter = filter;
self
}
pub fn with_pk_columns(mut self, pk_columns: Option<Vec<String>>) -> Self {
self.pk_columns = pk_columns;
self
}
#[inline]
fn find_batch(&self, row_pos: usize) -> Option<&BatchRange> {
let idx = self.batch_ranges.partition_point(|b| b.end <= row_pos);
self.batch_ranges
.get(idx)
.filter(|b| row_pos >= b.start && row_pos < b.end)
}
fn query_index(&self) -> Result<Vec<FtsHit>> {
let columns = self.query.columns();
if columns.len() > 1 {
return self.query_across_columns();
}
let Some((column, question)) = self.query.index_questions().into_iter().next() else {
return Err(Error::invalid_input(
"full-text search names no column to search".to_string(),
));
};
let index = self
.indexes
.index_answering(column, &question)
.ok_or_else(|| declined(column))?;
let query_expr = self.query.expr.clone();
let all_rows_visible = self.batch_ranges.last().is_none_or(|last| {
self.max_readable_row
.map(|max_readable| max_readable + 1 >= last.end as u64)
.unwrap_or(last.end == 0)
});
let pk_recency_is_noop = self.pk_columns.is_none()
|| (self.indexes.has_pk_index() && !self.indexes.pk_has_overrides());
let can_prune_in_index = self.filter.is_none() && pk_recency_is_noop && all_rows_visible;
let mut options = SearchOptions::new().with_include_tail(self.query.include_tail);
if can_prune_in_index {
options = options.with_wand_factor(self.query.wand_factor);
if let Some(limit) = self.query.limit {
options = options.with_limit(limit);
}
}
let query = FtsMemQuery {
expr: query_expr,
options,
granularity: self.query.document_granularity,
};
let Some(max_visible) = self.max_readable_row else {
return Ok(vec![]);
};
let Some(MemMatches::Ranked(ranked)) =
index.search(&query, &SearchContext::new(max_visible))?
else {
return Err(declined(column));
};
Ok(ranked
.into_iter()
.map(|hit| (hit.position, hit.element, hit.score))
.collect())
}
fn query_across_columns(&self) -> Result<Vec<FtsHit>> {
let Some(max_visible) = self.max_readable_row else {
return Ok(vec![]);
};
let options = SearchOptions::new().with_include_tail(self.query.include_tail);
let ctx = SearchContext::new(max_visible);
Ok(search_cross_column(&self.query.expr, |column, leaf| {
let query = FtsMemQuery {
expr: leaf.clone(),
options: options.clone(),
granularity: self.query.document_granularity,
};
let index = self
.indexes
.index_answering(column, &query)
.ok_or_else(|| declined(column))?;
match index.search(&query, &ctx)? {
Some(MemMatches::Ranked(ranked)) => {
Ok(ranked.into_iter().map(FtsEntry::from).collect())
}
_ => Err(declined(column)),
}
})?
.into_iter()
.map(|entry| (entry.row_position, entry.doc_index, entry.score))
.collect())
}
fn filter_by_visibility(&self, results: Vec<FtsHit>) -> Vec<FtsHit> {
let Some(max_readable) = self.max_readable_row else {
return vec![];
};
results
.into_iter()
.filter(|(pos, _, _)| *pos <= max_readable)
.collect()
}
fn materialize_rows_sorted(&self, results: &[FtsHit]) -> DataFusionResult<Vec<RecordBatch>> {
if results.is_empty() {
return Ok(vec![]);
}
let mut all_rows: Vec<u32> = Vec::with_capacity(results.len());
let mut all_scores: Vec<f32> = Vec::with_capacity(results.len());
let mut all_row_positions: Vec<u64> = Vec::with_capacity(results.len());
let mut all_doc_indices: Vec<Option<Vec<u32>>> = Vec::with_capacity(results.len());
let mut all_columns: Vec<Vec<Arc<dyn arrow_array::Array>>> = Vec::new();
let mut scan_batches = HashMap::<usize, RecordBatch>::new();
let first_batch = self.batch_store.get(0);
if let Some(stored) = first_batch {
for _ in 0..stored.data.num_columns() {
all_columns.push(Vec::with_capacity(results.len()));
}
}
for (pos, doc_index, score) in results {
let pos = *pos;
let score = *score;
if let Some(batch_range) = self.find_batch(pos as usize)
&& let Some(stored) = self.batch_store.get(batch_range.batch_id)
{
let data = match scan_batches.entry(batch_range.batch_id) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(scan_record_batch(&stored.data)?),
};
let row_in_batch = (pos as usize - batch_range.start) as u32;
let indices = UInt32Array::from(vec![row_in_batch]);
for (col_idx, col) in data.columns().iter().enumerate() {
let taken = arrow_select::take::take(col.as_ref(), &indices, None).unwrap();
if all_columns.len() <= col_idx {
all_columns.push(Vec::new());
}
all_columns[col_idx].push(taken);
}
all_rows.push(row_in_batch);
all_scores.push(score);
all_row_positions.push(pos);
all_doc_indices.push(doc_index.clone());
}
}
if all_scores.is_empty() {
return Ok(vec![]);
}
let mut final_columns: Vec<Arc<dyn arrow_array::Array>> = Vec::new();
for col_arrays in &all_columns {
if !col_arrays.is_empty() {
let refs: Vec<&dyn arrow_array::Array> =
col_arrays.iter().map(|a| a.as_ref()).collect();
let concatenated = arrow_select::concat::concat(&refs)?;
final_columns.push(concatenated);
}
}
let (final_columns, all_scores, all_row_positions, all_doc_indices) =
if let Some(ref predicate) = self.filter {
let Some(first) = self.batch_store.get(0) else {
return Ok(vec![]);
};
let data_batch =
RecordBatch::try_new(scan_record_batch(&first.data)?.schema(), final_columns)?;
let mask = predicate
.evaluate(&data_batch)?
.into_array(data_batch.num_rows())?;
let mask = mask
.as_any()
.downcast_ref::<BooleanArray>()
.ok_or_else(|| {
datafusion::error::DataFusionError::Internal(
"FTS prefilter predicate did not evaluate to boolean".to_string(),
)
})?;
let filtered_columns = data_batch
.columns()
.iter()
.map(|c| arrow_select::filter::filter(c.as_ref(), mask))
.collect::<std::result::Result<Vec<_>, _>>()?;
let filtered_scores: Vec<f32> = all_scores
.iter()
.zip(mask.iter())
.filter_map(|(s, keep)| keep.unwrap_or(false).then_some(*s))
.collect();
let filtered_positions: Vec<u64> = all_row_positions
.iter()
.zip(mask.iter())
.filter_map(|(p, keep)| keep.unwrap_or(false).then_some(*p))
.collect();
let filtered_doc_indices: Vec<Option<Vec<u32>>> = all_doc_indices
.iter()
.zip(mask.iter())
.filter(|(_, keep)| keep.unwrap_or(false))
.map(|(index, _)| index.clone())
.collect();
(
filtered_columns,
filtered_scores,
filtered_positions,
filtered_doc_indices,
)
} else {
(
final_columns,
all_scores,
all_row_positions,
all_doc_indices,
)
};
let (mut final_columns, mut all_scores, mut all_row_positions, mut all_doc_indices) = self
.filter_to_newest_pk(
final_columns,
all_scores,
all_row_positions,
all_doc_indices,
)?;
if all_scores.is_empty() {
return Ok(vec![]);
}
if let Some(limit) = self.query.limit
&& all_scores.len() > limit
{
final_columns = final_columns
.into_iter()
.map(|column| column.slice(0, limit))
.collect();
all_scores.truncate(limit);
all_row_positions.truncate(limit);
all_doc_indices.truncate(limit);
}
if self.with_doc_index {
let mut builder = ListBuilder::new(UInt32Builder::new()).with_field(Field::new(
"item",
DataType::UInt32,
false,
));
for doc_index in all_doc_indices {
let doc_index = doc_index.ok_or_else(|| {
datafusion::error::DataFusionError::Internal(
"element-document FTS result is missing its document coordinate"
.to_string(),
)
})?;
builder.values().append_slice(&doc_index);
builder.append(true);
}
final_columns.push(Arc::new(builder.finish()));
}
let n_hits = all_scores.len();
final_columns.push(Arc::new(Float32Array::from(all_scores)));
let mut projected_columns = if let Some(ref proj_indices) = self.projection {
let source_fields = self
.batch_store
.get(0)
.map(|stored| scan_record_batch(&stored.data))
.transpose()?
.map(|batch| batch.schema().fields().clone());
let mut projected: Vec<_> = match source_fields {
Some(fields) => take_projected_columns(
&final_columns,
&fields,
proj_indices,
self.output_schema.as_ref(),
n_hits,
)?,
None => proj_indices
.iter()
.map(|&i| final_columns[i].clone())
.collect(),
};
if self.with_doc_index {
projected.push(final_columns[final_columns.len() - 2].clone());
}
projected.push(final_columns.last().unwrap().clone());
projected
} else {
final_columns
};
if self.with_row_id {
projected_columns.push(Arc::new(UInt64Array::from(all_row_positions)));
}
let batch = RecordBatch::try_new(self.output_schema.clone(), projected_columns)?;
Ok(vec![batch])
}
fn filter_to_newest_pk(
&self,
final_columns: Vec<Arc<dyn arrow_array::Array>>,
all_scores: Vec<f32>,
all_row_positions: Vec<u64>,
all_doc_indices: Vec<Option<Vec<u32>>>,
) -> DataFusionResult<MaterializedFtsRows> {
let Some(pk_columns) = &self.pk_columns else {
return Ok((
final_columns,
all_scores,
all_row_positions,
all_doc_indices,
));
};
if pk_columns.is_empty() || all_scores.is_empty() {
return Ok((
final_columns,
all_scores,
all_row_positions,
all_doc_indices,
));
}
let Some(max_readable_row) = self.max_readable_row else {
return Ok((
final_columns,
all_scores,
all_row_positions,
all_doc_indices,
));
};
if self.indexes.has_pk_index() && !self.indexes.pk_has_overrides() {
return Ok((
final_columns,
all_scores,
all_row_positions,
all_doc_indices,
));
}
let Some(first) = self.batch_store.get(0) else {
return Ok((
final_columns,
all_scores,
all_row_positions,
all_doc_indices,
));
};
let newest_positions = if self.indexes.has_pk_index() {
None
} else {
Some(newest_pk_positions(
&self.batch_store,
pk_columns,
self.readable_count,
max_readable_row,
)?)
};
let data_batch =
RecordBatch::try_new(scan_record_batch(&first.data)?.schema(), final_columns)?;
let pk_indices = resolve_pk_indices(&data_batch, pk_columns)?;
let keep = (0..data_batch.num_rows())
.map(|row| {
Ok(match &newest_positions {
Some(newest) => newest.contains(&all_row_positions[row]),
None => {
let values: Vec<ScalarValue> = pk_indices
.iter()
.map(|&col| ScalarValue::try_from_array(data_batch.column(col), row))
.collect::<DataFusionResult<_>>()?;
self.indexes
.pk_is_newest(&values, all_row_positions[row], max_readable_row)
}
})
})
.collect::<DataFusionResult<Vec<_>>>()?;
let mask = BooleanArray::from_iter(keep.iter().copied());
let filtered_columns = data_batch
.columns()
.iter()
.map(|c| arrow_select::filter::filter(c.as_ref(), &mask))
.collect::<std::result::Result<Vec<_>, _>>()?;
let filtered_scores = all_scores
.into_iter()
.zip(keep.iter())
.filter_map(|(s, keep)| keep.then_some(s))
.collect();
let filtered_positions = all_row_positions
.into_iter()
.zip(keep.iter())
.filter_map(|(p, keep)| keep.then_some(p))
.collect();
let filtered_doc_indices = all_doc_indices
.into_iter()
.zip(keep.iter())
.filter_map(|(index, keep)| keep.then_some(index))
.collect();
Ok((
filtered_columns,
filtered_scores,
filtered_positions,
filtered_doc_indices,
))
}
}
impl DisplayAs for FtsIndexExec {
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result {
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
write!(
f,
"FtsIndexExec: columns={:?}, query_type={:?}, with_row_id={}",
self.query.columns(),
self.query.expr,
self.with_row_id
)
}
DisplayFormatType::TreeRender => {
write!(
f,
"FtsIndexExec\ncolumns={:?}\nquery_type={:?}\nwith_row_id={}",
self.query.columns(),
self.query.expr,
self.with_row_id
)
}
}
}
}
impl ExecutionPlan for FtsIndexExec {
fn name(&self) -> &str {
"FtsIndexExec"
}
fn schema(&self) -> SchemaRef {
self.output_schema.clone()
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
if !children.is_empty() {
return Err(datafusion::error::DataFusionError::Internal(
"FtsIndexExec does not have children".to_string(),
));
}
Ok(self)
}
fn execute(
&self,
_partition: usize,
_context: Arc<TaskContext>,
) -> DataFusionResult<SendableRecordBatchStream> {
let results = self.query_index()?;
let mut visible_results = self.filter_by_visibility(results);
visible_results.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
let batches = self.materialize_rows_sorted(&visible_results)?;
let stream = stream::iter(batches.into_iter().map(Ok)).boxed();
Ok(Box::pin(RecordBatchStreamAdapter::new(
self.output_schema.clone(),
stream,
)))
}
fn partition_statistics(&self, _partition: Option<usize>) -> DataFusionResult<Arc<Statistics>> {
Ok(Arc::new(Statistics {
num_rows: Precision::Absent,
total_byte_size: Precision::Absent,
column_statistics: Statistics::unknown_column(&self.schema()),
}))
}
fn metrics(&self) -> Option<MetricsSet> {
Some(self.metrics.clone_inner())
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
fn supports_limit_pushdown(&self) -> bool {
false
}
}
fn declined(column: &str) -> Error {
Error::internal(format!(
"the full-text index on '{column}' accepted a search, then did not answer it"
))
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::array::AsArray;
use arrow_array::builder::{ListBuilder, StringBuilder};
use arrow_array::{Array, Int32Array, StringArray};
use arrow_schema::{DataType, Field, Schema};
use futures::TryStreamExt;
use lance_index::scalar::InvertedIndexParams;
use lance_index::scalar::inverted::{DOC_INDEX_COL, DocumentGranularity};
fn create_test_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("text", DataType::Utf8, true),
]))
}
fn create_test_batch(schema: &Schema, start_id: i32) -> RecordBatch {
RecordBatch::try_new(
Arc::new(schema.clone()),
vec![
Arc::new(Int32Array::from(vec![start_id, start_id + 1, start_id + 2])),
Arc::new(StringArray::from(vec![
"hello world",
"goodbye world",
"hello again",
])),
],
)
.unwrap()
}
#[tokio::test]
async fn test_fts_index_search() {
let schema = create_test_schema();
let batch_store = Arc::new(BatchStore::with_capacity(100));
let mut registry = IndexStore::new();
registry.add_fts("text_idx".to_string(), 1, "text".to_string());
let batch = create_test_batch(&schema, 0);
registry.insert(&batch, 0).unwrap();
batch_store.append(batch).unwrap();
let indexes = Arc::new(registry);
let query = FtsQuery::match_query("text", "hello");
let exec = FtsIndexExec::new(batch_store, indexes, query, 1, None, schema, false).unwrap();
let ctx = Arc::new(TaskContext::default());
let stream = exec.execute(0, ctx).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total_rows, 2);
let result_schema = batches[0].schema();
assert!(result_schema.field_with_name(SCORE_COLUMN).is_ok());
}
#[tokio::test]
async fn test_element_document_fts_index_search() {
let mut tags = ListBuilder::new(StringBuilder::new());
tags.values().append_value("alpha beta");
tags.values().append_value("beta gamma");
tags.append(true);
tags.values().append_value("beta");
tags.append(true);
let tags = tags.finish();
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("tags", tags.data_type().clone(), true),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![0, 1])), Arc::new(tags)],
)
.unwrap();
let batch_store = Arc::new(BatchStore::with_capacity(10));
let mut registry = IndexStore::new();
registry
.add_fts_with_params(
"tags_element_idx".to_string(),
1,
"tags".to_string(),
InvertedIndexParams::default()
.document_granularity(DocumentGranularity::ListElement),
)
.unwrap();
registry.insert(&batch, 0).unwrap();
batch_store.append(batch).unwrap();
let exec = FtsIndexExec::new(
batch_store,
Arc::new(registry),
FtsQuery::match_query("tags", "beta")
.with_document_granularity(DocumentGranularity::ListElement),
1,
None,
schema,
false,
)
.unwrap();
let stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let mut hits = Vec::new();
for batch in batches {
let ids = batch["id"].as_primitive::<arrow::datatypes::Int32Type>();
let coordinates = batch[DOC_INDEX_COL].as_list::<i32>();
for row in 0..batch.num_rows() {
let coordinate = coordinates
.value(row)
.as_primitive::<arrow::datatypes::UInt32Type>()
.value(0);
hits.push((ids.value(row), coordinate));
}
}
hits.sort_unstable();
assert_eq!(hits, vec![(0, 0), (0, 1), (1, 0)]);
}
#[tokio::test]
async fn test_fts_index_visibility() {
let schema = create_test_schema();
let batch_store = Arc::new(BatchStore::with_capacity(100));
let mut registry = IndexStore::new();
registry.add_fts("text_idx".to_string(), 1, "text".to_string());
let batch1 = create_test_batch(&schema, 0);
let batch2 = create_test_batch(&schema, 5);
registry.insert(&batch1, 0).unwrap();
registry.insert(&batch2, 3).unwrap(); batch_store.append(batch1).unwrap();
batch_store.append(batch2).unwrap();
let indexes = Arc::new(registry);
let query = FtsQuery::match_query("text", "hello");
let exec = FtsIndexExec::new(
batch_store.clone(),
indexes.clone(),
query.clone(),
1,
None,
schema.clone(),
false,
)
.unwrap();
let ctx = Arc::new(TaskContext::default());
let stream = exec.execute(0, ctx).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total_rows, 2);
let exec = FtsIndexExec::new(batch_store, indexes, query, 2, None, schema, false).unwrap();
let ctx = Arc::new(TaskContext::default());
let stream = exec.execute(0, ctx).unwrap();
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total_rows, 4); }
#[test]
fn test_score_column_name() {
assert_eq!(SCORE_COLUMN, "_score");
}
}