lance_nd/
scan.rs

1//! Read a Lance table as a stream of nd record batches.
2
3use arrow::datatypes::{Field, Schema};
4use futures::stream::{self, BoxStream};
5use futures::{StreamExt, TryStreamExt};
6use lance::Dataset;
7use lance::dataset::scanner::{MaterializationStyle, Scanner};
8use lance::deps::datafusion::logical_expr::Expr;
9use lance::table::format::Fragment;
10use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu};
11use nd_arrow_array::NdRecordBatch;
12use nd_arrow_array::extension::GLOBAL_ATTRIBUTES_FIELD;
13
14use crate::encode::{RowDecoder, variable_field};
15use crate::error::{Error, Result};
16use crate::format::check_format;
17use crate::prune::prune_filter;
18use crate::stats::is_stats_name;
19
20/// The variables of the table, as logical fields typed by their values.
21pub fn variables(dataset: &Dataset) -> Result<Vec<Field>> {
22    check_format(dataset)?;
23    Schema::from(dataset.schema())
24        .fields()
25        .iter()
26        .filter(|f| f.name() != GLOBAL_ATTRIBUTES_FIELD && !is_stats_name(f.name()))
27        .map(|f| variable_field(f))
28        .collect()
29}
30
31/// A scan over the rows of the table.
32pub fn scan(dataset: &Dataset) -> NdScanner<'_> {
33    NdScanner {
34        dataset,
35        projection: None,
36        filter: None,
37        filter_expr: None,
38        lance_filter: None,
39        fragments: None,
40        limit: None,
41        batch_size: DEFAULT_BATCH_SIZE,
42        batch_size_bytes: DEFAULT_BATCH_SIZE_BYTES,
43        batch_readahead: None,
44        decode_parallelism: None,
45    }
46}
47
48/// Default rows in one Lance batch.
49const DEFAULT_BATCH_SIZE: usize = 256;
50
51/// Default bytes in one Lance batch. A large row gets a batch of its own.
52const DEFAULT_BATCH_SIZE_BYTES: u64 = 32 << 20;
53
54/// A scan that yields one nd record batch per row, in table order.
55///
56/// Lance reads batches of rows, limited by rows and by bytes. Several
57/// batches decode to nd record batches at the same time.
58pub struct NdScanner<'a> {
59    dataset: &'a Dataset,
60    projection: Option<Vec<String>>,
61    filter: Option<String>,
62    filter_expr: Option<Expr>,
63    lance_filter: Option<Expr>,
64    fragments: Option<Vec<Fragment>>,
65    limit: Option<i64>,
66    batch_size: usize,
67    batch_size_bytes: u64,
68    batch_readahead: Option<usize>,
69    decode_parallelism: Option<usize>,
70}
71
72impl NdScanner<'_> {
73    /// Read only these variables. A row keeps its global attributes.
74    /// A row without any of these variables gives a batch without columns,
75    /// so an empty list reads only the global attributes.
76    ///
77    /// Each name must be a whole variable. A nested field, such as
78    /// `sst.attrs.units`, gives [`Error::InvalidArgument`] in the scan.
79    pub fn project(mut self, variables: &[&str]) -> Self {
80        self.projection = Some(variables.iter().map(|v| v.to_string()).collect());
81        self
82    }
83
84    /// Read only the rows that match a Lance SQL filter.
85    ///
86    /// A scan cannot have both this filter and a DataFusion filter, see
87    /// [`NdScanner::filter_expr`] and [`NdScanner::lance_filter`].
88    pub fn filter(mut self, filter: impl Into<String>) -> Self {
89        self.filter = Some(filter.into());
90        self
91    }
92
93    /// Skip the rows that surely have no value that matches `expr`.
94    ///
95    /// `expr` is a DataFusion filter on the table columns, for example
96    /// `get_field(col("lat"), "values").between(lit(17.5), lit(25.4))`. A
97    /// part on `<variable>.values` uses the stats of the table, see
98    /// [`crate::prune_filter`]. A kept row comes back whole, also with values
99    /// outside the filter. A table without stats keeps all rows for a value
100    /// part. Each other part, for example on attributes, filters the rows
101    /// exactly.
102    pub fn filter_expr(mut self, expr: Expr) -> Self {
103        self.filter_expr = Some(expr);
104        self
105    }
106
107    /// Read only the rows that match `expr`, a DataFusion filter on the
108    /// Lance table columns.
109    ///
110    /// Lance gets `expr` with no change, so `expr` must use the table
111    /// columns: the nd columns, `__global_attrs__` and the stats columns. With
112    /// [`NdScanner::filter_expr`], a row must match both filters.
113    pub fn lance_filter(mut self, expr: Expr) -> Self {
114        self.lance_filter = Some(expr);
115        self
116    }
117
118    /// Read only the rows of these fragments, in the order given.
119    pub fn fragments(mut self, fragments: Vec<Fragment>) -> Self {
120        self.fragments = Some(fragments);
121        self
122    }
123
124    /// Read at most `limit` rows.
125    pub fn limit(mut self, limit: usize) -> Self {
126        self.limit = Some(limit as i64);
127        self
128    }
129
130    /// The most rows in one Lance batch. The default is 256.
131    pub fn batch_size(mut self, rows: usize) -> Self {
132        self.batch_size = rows.max(1);
133        self
134    }
135
136    /// The approximate most bytes in one Lance batch. The default is 32 MiB.
137    /// A row larger than this is a batch of its own.
138    pub fn batch_size_bytes(mut self, bytes: u64) -> Self {
139        self.batch_size_bytes = bytes.max(1);
140        self
141    }
142
143    /// The most Lance batches that decode at the same time. The default is
144    /// the number of compute CPUs of Lance.
145    pub fn decode_parallelism(mut self, batches: usize) -> Self {
146        self.decode_parallelism = Some(batches.max(1));
147        self
148    }
149
150    /// The most Lance batches that Lance decodes ahead of the consumer.
151    pub fn batch_readahead(mut self, batches: usize) -> Self {
152        self.batch_readahead = Some(batches);
153        self
154    }
155
156    /// The Lance plan of the scan, for example to see which columns the
157    /// filter read loads.
158    pub async fn explain_plan(&self, verbose: bool) -> Result<String> {
159        let (scanner, _) = self.lance_scanner()?;
160        Ok(scanner.explain_plan(verbose).await?)
161    }
162
163    /// A stream of the selected rows, in table order.
164    #[tracing::instrument(level = "debug", skip_all, fields(uri = %self.dataset.uri()), err)]
165    pub async fn try_into_stream(self) -> Result<BoxStream<'static, Result<NdRecordBatch>>> {
166        let (scanner, rows_only) = self.lance_scanner()?;
167        let parallelism = self
168            .decode_parallelism
169            .unwrap_or_else(get_num_compute_intensive_cpus);
170        let batches = scanner.try_into_stream().await?;
171        // Decode on the Lance CPU pool. `buffered` keeps the table order.
172        Ok(batches
173            .map(move |batch| async move {
174                let batch = match rows_only {
175                    true => batch?.project(&[])?,
176                    false => batch?,
177                };
178                spawn_cpu(move || {
179                    RowDecoder::try_new(&batch)?
180                        .into_rows()
181                        .collect::<Result<Vec<_>>>()
182                })
183                .await
184            })
185            .buffered(parallelism)
186            .map_ok(|rows| stream::iter(rows.into_iter().map(Ok)))
187            .try_flatten()
188            .boxed())
189    }
190
191    /// The Lance scanner, and `true` when it reads only row ids.
192    fn lance_scanner(&self) -> Result<(Scanner, bool)> {
193        check_format(self.dataset)?;
194        let mut scanner = self.dataset.scan();
195        scanner.batch_size(self.batch_size);
196        scanner.batch_size_bytes(self.batch_size_bytes);
197        let mut rows_only = false;
198        if let Some(projection) = &self.projection {
199            let schema = self.dataset.schema();
200            let is_column = |name: &str| schema.fields.iter().any(|f| f.name == name);
201            if let Some(unknown) = projection
202                .iter()
203                .find(|v| *v == GLOBAL_ATTRIBUTES_FIELD || is_stats_name(v) || !is_column(v))
204            {
205                return Err(Error::InvalidArgument(format!(
206                    "'{unknown}' is not a variable of the table: a projection takes whole variables, not nested fields"
207                )));
208            }
209            let mut columns = projection.clone();
210            if is_column(GLOBAL_ATTRIBUTES_FIELD) {
211                columns.push(GLOBAL_ATTRIBUTES_FIELD.to_string());
212            }
213            // Lance needs one column, so read the row ids and drop them.
214            rows_only = columns.is_empty();
215            if rows_only {
216                scanner.with_row_id();
217            }
218            scanner.project(&columns)?;
219        }
220        if self.filter.is_some() && (self.filter_expr.is_some() || self.lance_filter.is_some()) {
221            return Err(Error::InvalidArgument(
222                "a scan cannot have both filter and a DataFusion filter".to_string(),
223            ));
224        }
225        let mut filtered = false;
226        if let Some(filter) = &self.filter {
227            scanner.filter(filter)?;
228            filtered = true;
229        }
230        let pruned = match &self.filter_expr {
231            Some(expr) => prune_filter(expr, &Schema::from(self.dataset.schema()))?,
232            None => None,
233        };
234        if let Some(expr) = pruned
235            .into_iter()
236            .chain(self.lance_filter.clone())
237            .reduce(Expr::and)
238        {
239            scanner.filter_expr(expr);
240            filtered = true;
241        }
242        // Read the filter columns first, then the variables of the matching
243        // rows only. Lance reads a struct of lists early by default.
244        if filtered {
245            scanner.materialization_style(MaterializationStyle::AllLate);
246        }
247        if let Some(fragments) = &self.fragments {
248            scanner.with_fragments(fragments.clone());
249        }
250        if self.limit.is_some() {
251            scanner.limit(self.limit, None)?;
252        }
253        if let Some(rows) = self.batch_readahead {
254            scanner.batch_readahead(rows);
255        }
256        Ok((scanner, rows_only))
257    }
258
259    /// Collect the selected rows.
260    pub async fn try_collect(self) -> Result<Vec<NdRecordBatch>> {
261        self.try_into_stream().await?.try_collect().await
262    }
263}
lance_nd/scan.rs

lance_nd/
scan.rs
