feat: add initial scanner statistics

python/python/lance/dataset.py

@@ -17,6 +17,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Dict,
     Iterable,
     Iterator,
@@ -43,6 +44,7 @@
 from .fragment import FragmentMetadata, LanceFragment
 from .lance import (
     CleanupStats,
+    LanceScanStats,
     _Dataset,
     _MergeInsertBuilder,
     _Operation,
@@ -279,6 +281,7 @@ def scanner(
         io_buffer_size: Optional[int] = None,
         late_materialization: Optional[bool | List[str]] = None,
         use_scalar_index: Optional[bool] = None,
+        stats_handler: Optional[Union[Callable[[LanceScanStats], None], str]] = None,
     ) -> LanceScanner:
         """Return a Scanner that can support various pushdowns.
 
@@ -367,6 +370,17 @@ def scanner(
         fast_search:  bool, default False
             If True, then the search will only be performed on the indexed data, which
             yields faster search time.
+        stats_handler: 'full', 'brief', or Callable, default None
+            If None then stats will not be collected.
+
+            If 'full' then detailed stats will be written to the log file (with DEBUG
+            level).  This includes the scan plan.
+
+            If 'brief' then brief stats (overall timing) will be written to the log
+            file.
+
+            If a callable, then the callable will be provided an instance of
+            LanceScanStats when the scan is complete.
 
         Notes
         -----
@@ -415,6 +429,7 @@ def setopt(opt, val):
         setopt(builder.use_stats, use_stats)
         setopt(builder.use_scalar_index, use_scalar_index)
         setopt(builder.fast_search, fast_search)
+        setopt(builder.stats_handler, stats_handler)
 
         # columns=None has a special meaning. we can't treat it as "user didn't specify"
         if self._default_scan_options is None:
@@ -488,6 +503,7 @@ def to_table(
         io_buffer_size: Optional[int] = None,
         late_materialization: Optional[bool | List[str]] = None,
         use_scalar_index: Optional[bool] = None,
+        stats_handler: Optional[Union[Callable[[LanceScanStats], None], str]] = None,
     ) -> pa.Table:
         """Read the data into memory as a pyarrow Table.
 
@@ -555,6 +571,17 @@ def to_table(
                 currently only supports a single column in the columns list.
             - query: str
                 The query string to search for.
+        stats_handler: 'full', 'brief', or Callable, default None
+            If None then stats will not be collected.
+
+            If 'full' then detailed stats will be written to the log file (with DEBUG
+            level).  This includes the scan plan.
+
+            If 'brief' then brief stats (overall timing) will be written to the log
+            file.
+
+            If a callable, then the callable will be provided an instance of
+            LanceScanStats when the scan is complete.
 
         Notes
         -----
@@ -581,6 +608,7 @@ def to_table(
             use_stats=use_stats,
             fast_search=fast_search,
             full_text_query=full_text_query,
+            stats_handler=stats_handler,
         ).to_table()
 
     @property
@@ -635,6 +663,7 @@ def to_batches(
         io_buffer_size: Optional[int] = None,
         late_materialization: Optional[bool | List[str]] = None,
         use_scalar_index: Optional[bool] = None,
+        stats_handler: Optional[Union[Callable[[LanceScanStats], None], str]] = None,
         **kwargs,
     ) -> Iterator[pa.RecordBatch]:
         """Read the dataset as materialized record batches.
@@ -666,6 +695,7 @@ def to_batches(
             with_row_address=with_row_address,
             use_stats=use_stats,
             full_text_query=full_text_query,
+            stats_handler=stats_handler,
         ).to_batches()
 
     def sample(
@@ -2590,6 +2620,7 @@ def __init__(self, ds: LanceDataset):
         self._fast_search = False
         self._full_text_query = None
         self._use_scalar_index = None
+        self._stats_handler = None
 
     def apply_defaults(self, default_opts: Dict[str, Any]) -> ScannerBuilder:
         for key, value in default_opts.items():
@@ -2864,6 +2895,28 @@ def full_text_search(
         self._full_text_query = {"query": query, "columns": columns}
         return self
 
+    def stats_handler(
+        self, handler: Union[Callable[[LanceScanStats], None], str]
+    ) -> ScannerBuilder:
+        """
+        Sets the handler for scan statistics
+
+        Scan statistics are gathered while the scan is run and, once the scan is
+        complete, these statistics are given to the handler.
+
+        If 'full' then full stats (including the plan) are collected and logged.
+
+        If 'brief' then only global scan timers are collected and logged.
+
+        If the handler is a callable then the stats will be provided to the callable.
+
+        Note: the format of the log message provided by 'full' and 'brief' is NOT
+        stable and subject to change.  Prefer using a callable instead of parsing the
+        log messages.
+        """
+        self._stats_handler = handler
+        return self
+
     def to_scanner(self) -> LanceScanner:
         scanner = self.ds._ds.scanner(
             self._columns,
@@ -2887,6 +2940,7 @@ def to_scanner(self) -> LanceScanner:
             self._full_text_query,
             self._late_materialization,
             self._use_scalar_index,
+            self._stats_handler,
         )
         return LanceScanner(scanner, self.ds)
 

python/python/lance/lance/__init__.pyi

@@ -93,3 +93,12 @@ class LanceBlobFile:
     def size(self) -> int: ...
     def readall(self) -> bytes: ...
     def readinto(self, b: bytearray) -> int: ...
+
+class LanceScanStats:
+    start: int
+    end: int
+    wall_clock_duration: float
+    wall_clock_throughput: float
+    output_rows: int
+    estimated_output_bytes: int
+    plan: Optional[str]

python/python/tests/test_dataset.py

@@ -2095,6 +2095,31 @@ def test_dataset_restore(tmp_path: Path):
     assert dataset.count_rows() == 100
 
 
+def test_scan_stats(tmp_path: Path):
+    data = pa.table({"a": range(100)})
+    dataset = lance.write_dataset(data, tmp_path)
+
+    dataset.to_table(stats_handler="full")
+    dataset.to_table(stats_handler="brief")
+
+    now = datetime.now()
+    unix_now = time.mktime(now.timetuple())
+
+    def assert_stats(stats: lance.lance.LanceScanStats):
+        # Start/end should be within 5 minutes of now
+        assert stats.start / 1000 < unix_now + 300
+        assert stats.start / 1000 > unix_now - 300
+        assert stats.end / 1000 < unix_now + 300
+        assert stats.end / 1000 > unix_now - 300
+        assert stats.wall_clock_duration > 0
+        assert stats.wall_clock_throughput > 0
+        assert stats.output_rows == 100
+        assert stats.estimated_output_bytes == 800
+        assert stats.plan is not None
+
+    dataset.to_table(stats_handler=assert_stats)
+
+
 def test_mixed_mode_overwrite(tmp_path: Path):
     data = pa.table({"a": range(100)})
     dataset = lance.write_dataset(data, tmp_path, data_storage_version="legacy")

python/src/dataset.rs

@@ -29,6 +29,7 @@ use arrow_array::Array;
 use futures::{StreamExt, TryFutureExt};
 use lance::dataset::builder::DatasetBuilder;
 use lance::dataset::refs::{Ref, TagContents};
+use lance::dataset::scanner::stats::ScanStatisticsHandler;
 use lance::dataset::scanner::MaterializationStyle;
 use lance::dataset::transaction::{
     validate_operation, RewriteGroup as LanceRewriteGroup, RewrittenIndex as LanceRewrittenIndex,
@@ -63,7 +64,7 @@ use lance_table::io::commit::CommitHandler;
 use object_store::path::Path;
 use pyo3::exceptions::{PyStopIteration, PyTypeError};
 use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyInt, PyList, PySet, PyString};
+use pyo3::types::{PyBytes, PyInt, PyList, PySet, PyString, PyTuple};
 use pyo3::{
     exceptions::{PyIOError, PyKeyError, PyValueError},
     pyclass,
@@ -75,6 +76,7 @@ use uuid::Uuid;
 
 use crate::error::PythonErrorExt;
 use crate::fragment::{FileFragment, FragmentMetadata};
+use crate::scanner::LanceScanStats;
 use crate::schema::LanceSchema;
 use crate::session::Session;
 use crate::RT;
@@ -603,6 +605,7 @@ impl Dataset {
         full_text_query: Option<&PyDict>,
         late_materialization: Option<PyObject>,
         use_scalar_index: Option<bool>,
+        stats_handler: Option<Bound<PyAny>>,
     ) -> PyResult<Scanner> {
         let mut scanner: LanceScanner = self_.ds.scan();
         match (columns, columns_with_transform) {
@@ -682,6 +685,37 @@ impl Dataset {
             scanner.fragment_readahead(fragment_readahead);
         }
 
+        if let Some(stats_handler) = stats_handler {
+            if stats_handler.downcast::<PyString>().is_ok() {
+                let stats_handler = stats_handler.extract::<String>()?;
+                scanner.stats_handler(
+                    stats_handler
+                        .parse::<ScanStatisticsHandler>()
+                        .infer_error()?,
+                );
+            } else if stats_handler.is_callable() {
+                let stats_handler = stats_handler.unbind();
+                let stats_handler = ScanStatisticsHandler::Custom(Arc::new(move |stats| {
+                    let wrapped_stats = LanceScanStats::new(stats);
+                    let stats_handler = stats_handler.clone();
+                    Python::with_gil(move |py| {
+                        let args = PyTuple::new(py, vec![wrapped_stats.into_py(py)]);
+                        stats_handler.call1(py, args)
+                    })
+                    .map_err(|err| lance_core::Error::Wrapped {
+                        error: err.into(),
+                        location: location!(),
+                    })?;
+                    Ok(())
+                }));
+                scanner.stats_handler(stats_handler);
+            } else {
+                return Err(PyValueError::new_err(
+                    "stats_handler must be a string or a callable",
+                ));
+            }
+        }
+
         scanner.scan_in_order(scan_in_order.unwrap_or(true));
 
         if with_row_id.unwrap_or(false) {

python/src/lib.rs

@@ -44,6 +44,7 @@ use futures::StreamExt;
 use lance_index::DatasetIndexExt;
 use pyo3::exceptions::{PyIOError, PyValueError};
 use pyo3::prelude::*;
+use scanner::LanceScanStats;
 use session::Session;
 
 #[macro_use]
@@ -123,6 +124,7 @@ fn lance(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<LanceColumnMetadata>()?;
     m.add_class::<LancePageMetadata>()?;
     m.add_class::<LanceBufferDescriptor>()?;
+    m.add_class::<LanceScanStats>()?;
     m.add_class::<DataFile>()?;
     m.add_class::<BFloat16>()?;
     m.add_class::<CleanupStats>()?;