feat: make the index cache size (in bytes) available

Cargo.toml

@@ -99,6 +99,7 @@ datafusion-execution = "37.1"
 datafusion-physical-expr = { version = "37.1", features = [
     "regex_expressions",
 ] }
+deepsize = "0.2.0"
 either = "1.0"
 futures = "0.3"
 http = "0.2.9"

python/python/lance/dataset.py

@@ -51,6 +51,7 @@
 )
 from .lance import CompactionMetrics as CompactionMetrics
 from .lance import __version__ as __version__
+from .lance import _Session as Session
 from .optimize import Compaction
 from .schema import LanceSchema
 from .util import td_to_micros
@@ -1517,6 +1518,12 @@ def create_index(
         self._ds.create_index(column, index_type, name, replace, kwargs)
         return LanceDataset(self.uri, index_cache_size=index_cache_size)
 
+    def session(self) -> Session:
+        """
+        Return the dataset session, which holds the dataset's state.
+        """
+        return self._ds.session()
+
     @staticmethod
     def _commit(
         base_uri: Union[str, Path],

python/python/lance/lance/__init__.pyi

@@ -76,3 +76,6 @@ class LanceFileMetadata:
     num_global_buffer_bytes: int
     global_buffers: List[LanceBufferDescriptor]
     columns: List[LanceColumnMetadata]
+
+class _Session:
+    def size_bytes(self) -> int: ...

python/python/tests/test_session.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+
+from pathlib import Path
+
+import lance
+import pyarrow as pa
+
+
+def test_cache_size_bytes(
+    tmp_path: Path,
+):
+    data = pa.table({"a": range(1000)})
+    lance.write_dataset(data, tmp_path, max_rows_per_file=250)
+
+    ds = lance.dataset(tmp_path)
+
+    initial_size = ds.session().size_bytes()
+
+    ds.scanner().to_table()
+
+    after_scan_size = ds.session().size_bytes()
+
+    assert after_scan_size > initial_size

python/src/dataset.rs

@@ -65,6 +65,7 @@ use snafu::{location, Location};
 
 use crate::fragment::{FileFragment, FragmentMetadata};
 use crate::schema::LanceSchema;
+use crate::session::Session;
 use crate::RT;
 use crate::{LanceReader, Scanner};
 
@@ -972,6 +973,10 @@ impl Dataset {
         Ok(self.ds.index_cache_hit_rate())
     }
 
+    fn session(&self) -> Session {
+        Session::new(self.ds.session())
+    }
+
     #[staticmethod]
     fn commit(
         dataset_uri: &str,

python/src/lib.rs

@@ -43,6 +43,7 @@ use futures::StreamExt;
 use lance_index::DatasetIndexExt;
 use pyo3::exceptions::{PyIOError, PyValueError};
 use pyo3::prelude::*;
+use session::Session;
 
 #[macro_use]
 extern crate lazy_static;
@@ -59,6 +60,7 @@ pub(crate) mod fragment;
 pub(crate) mod reader;
 pub(crate) mod scanner;
 pub(crate) mod schema;
+pub(crate) mod session;
 pub(crate) mod tracing;
 pub(crate) mod updater;
 pub(crate) mod utils;
@@ -127,6 +129,7 @@ fn lance(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyCompactionPlan>()?;
     m.add_class::<PyRewriteResult>()?;
     m.add_class::<PyCompactionMetrics>()?;
+    m.add_class::<Session>()?;
     m.add_class::<TraceGuard>()?;
     m.add_class::<schema::LanceSchema>()?;
     m.add_wrapped(wrap_pyfunction!(bfloat16_array))?;

python/src/session.rs

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use std::sync::Arc;
+
+use pyo3::{pyclass, pymethods};
+
+use lance::session::Session as LanceSession;
+
+/// The Session holds stateful information for a dataset.
+///
+/// The session contains caches for opened indices and file metadata.
+#[pyclass(name = "_Session", module = "_lib")]
+#[derive(Clone)]
+pub struct Session {
+    inner: Arc<LanceSession>,
+}
+
+impl Session {
+    pub fn new(inner: Arc<LanceSession>) -> Self {
+        Self { inner }
+    }
+}
+
+#[pymethods]
+impl Session {
+    /// Return the current size of the session in bytes
+    pub fn size_bytes(&self) -> u64 {
+        self.inner.size_bytes()
+    }
+}

rust/lance-core/Cargo.toml

@@ -22,6 +22,7 @@ bytes.workspace = true
 chrono.workspace = true
 datafusion-common = { workspace = true, optional = true }
 datafusion-sql = { workspace = true, optional = true }
+deepsize.workspace = true
 futures.workspace = true
 lazy_static.workspace = true
 mock_instant.workspace = true

rust/lance-core/src/cache.rs

@@ -6,6 +6,7 @@
 use std::any::{Any, TypeId};
 use std::sync::Arc;
 
+use deepsize::{Context, DeepSizeOf};
 use futures::Future;
 use moka::sync::Cache;
 use object_store::path::Path;
@@ -17,12 +18,46 @@ pub const DEFAULT_METADATA_CACHE_SIZE: usize = 128;
 
 type ArcAny = Arc<dyn Any + Send + Sync>;
 
+#[derive(Clone)]
+struct SizedRecord {
+    record: ArcAny,
+    size_accessor: Arc<dyn Fn(ArcAny) -> usize + Send + Sync>,
+}
+
+impl std::fmt::Debug for SizedRecord {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SizedRecord")
+            .field("record", &self.record)
+            .finish()
+    }
+}
+
+impl SizedRecord {
+    fn new<T: DeepSizeOf + Send + Sync + 'static>(record: Arc<T>) -> Self {
+        let size_accessor =
+            |record: ArcAny| -> usize { record.downcast_ref::<T>().unwrap().deep_size_of() };
+        Self {
+            record,
+            size_accessor: Arc::new(size_accessor),
+        }
+    }
+}
+
 /// Cache for various metadata about files.
 ///
 /// The cache is keyed by the file path and the type of metadata.
 #[derive(Clone, Debug)]
 pub struct FileMetadataCache {
-    cache: Arc<Cache<(Path, TypeId), ArcAny>>,
+    cache: Arc<Cache<(Path, TypeId), SizedRecord>>,
+}
+
+impl DeepSizeOf for FileMetadataCache {
+    fn deep_size_of_children(&self, _: &mut Context) -> usize {
+        self.cache
+            .iter()
+            .map(|(_, v)| (v.size_accessor)(v.record))
+            .sum()
+    }
 }
 
 impl FileMetadataCache {
@@ -35,19 +70,20 @@ impl FileMetadataCache {
     pub fn get<T: Send + Sync + 'static>(&self, path: &Path) -> Option<Arc<T>> {
         self.cache
             .get(&(path.to_owned(), TypeId::of::<T>()))
-            .map(|metadata| metadata.clone().downcast::<T>().unwrap())
+            .map(|metadata| metadata.record.clone().downcast::<T>().unwrap())
     }
 
-    pub fn insert<T: Send + Sync + 'static>(&self, path: Path, metadata: Arc<T>) {
-        self.cache.insert((path, TypeId::of::<T>()), metadata);
+    pub fn insert<T: DeepSizeOf + Send + Sync + 'static>(&self, path: Path, metadata: Arc<T>) {
+        self.cache
+            .insert((path, TypeId::of::<T>()), SizedRecord::new(metadata));
     }
 
     /// Get an item
     ///
     /// If it exists in the cache return that
     ///
     /// If it doesn't then run `loader` to load the item, insert into cache, and return
-    pub async fn get_or_insert<T: Send + Sync + 'static, F, Fut>(
+    pub async fn get_or_insert<T: DeepSizeOf + Send + Sync + 'static, F, Fut>(
         &self,
         path: &Path,
         loader: F,

rust/lance-core/src/datatypes.rs

@@ -8,6 +8,7 @@ use std::sync::Arc;
 
 use arrow_array::ArrayRef;
 use arrow_schema::{DataType, Field as ArrowField, TimeUnit};
+use deepsize::DeepSizeOf;
 use lance_arrow::bfloat16::{
     is_bfloat16_field, ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY, BFLOAT16_EXT_NAME,
 };
@@ -24,7 +25,7 @@ pub use schema::Schema;
 
 /// LogicalType is a string presentation of arrow type.
 /// to be serialized into protobuf.
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, DeepSizeOf)]
 pub struct LogicalType(String);
 
 impl fmt::Display for LogicalType {
@@ -324,6 +325,15 @@ pub struct Dictionary {
     pub values: Option<ArrayRef>,
 }
 
+impl DeepSizeOf for Dictionary {
+    fn deep_size_of_children(&self, _context: &mut deepsize::Context) -> usize {
+        self.values
+            .as_ref()
+            .map(|v| v.get_array_memory_size())
+            .unwrap_or(0)
+    }
+}
+
 impl PartialEq for Dictionary {
     fn eq(&self, other: &Self) -> bool {
         match (&self.values, &other.values) {

rust/lance-core/src/datatypes/field.rs

@@ -18,6 +18,7 @@ use arrow_array::{
     ArrayRef,
 };
 use arrow_schema::{DataType, Field as ArrowField};
+use deepsize::DeepSizeOf;
 use lance_arrow::{bfloat16::ARROW_EXT_NAME_KEY, *};
 use snafu::{location, Location};
 
@@ -34,7 +35,7 @@ pub struct SchemaCompareOptions {
     pub compare_field_ids: bool,
 }
 /// Encoding enum.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
 pub enum Encoding {
     /// Plain encoding.
     Plain,
@@ -48,7 +49,7 @@ pub enum Encoding {
 
 /// Lance Schema Field
 ///
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, DeepSizeOf)]
 pub struct Field {
     pub name: String,
     pub id: i32,

rust/lance-core/src/datatypes/schema.rs

@@ -10,14 +10,15 @@ use std::{
 
 use arrow_array::RecordBatch;
 use arrow_schema::{Field as ArrowField, Schema as ArrowSchema};
+use deepsize::DeepSizeOf;
 use lance_arrow::*;
 use snafu::{location, Location};
 
 use super::field::{Field, SchemaCompareOptions};
 use crate::{Error, Result};
 
 /// Lance Schema.
-#[derive(Default, Debug, Clone)]
+#[derive(Default, Debug, Clone, DeepSizeOf)]
 pub struct Schema {
     /// Top-level fields in the dataset.
     pub fields: Vec<Field>,

rust/lance-core/src/utils/deletion.rs

@@ -4,6 +4,7 @@
 use std::{collections::HashSet, ops::Range};
 
 use arrow_array::BooleanArray;
+use deepsize::{Context, DeepSizeOf};
 use roaring::RoaringBitmap;
 
 /// Threshold for when a DeletionVector::Set should be promoted to a DeletionVector::Bitmap.
@@ -18,6 +19,17 @@ pub enum DeletionVector {
     Bitmap(RoaringBitmap),
 }
 
+impl DeepSizeOf for DeletionVector {
+    fn deep_size_of_children(&self, context: &mut Context) -> usize {
+        match self {
+            Self::NoDeletions => 0,
+            Self::Set(set) => set.deep_size_of_children(context),
+            // Inexact but probably close enough
+            Self::Bitmap(bitmap) => bitmap.serialized_size(),
+        }
+    }
+}
+
 impl DeletionVector {
     #[allow(dead_code)] // Used in tests
     pub fn len(&self) -> usize {

rust/lance-file/Cargo.toml

@@ -27,6 +27,7 @@ async-trait.workspace = true
 byteorder.workspace = true
 bytes.workspace = true
 datafusion-common.workspace = true
+deepsize.workspace = true
 futures.workspace = true
 lance-datagen.workspace = true
 log.workspace = true

rust/lance-file/src/format/metadata.rs

@@ -6,12 +6,13 @@ use std::ops::Range;
 
 use crate::datatypes::{Fields, FieldsWithMeta};
 use crate::format::pb;
+use deepsize::DeepSizeOf;
 use lance_core::datatypes::Schema;
 use lance_core::{Error, Result};
 use lance_io::traits::ProtoStruct;
 use snafu::{location, Location};
 /// Data File Metadata
-#[derive(Debug, Default, PartialEq)]
+#[derive(Debug, Default, DeepSizeOf, PartialEq)]
 pub struct Metadata {
     /// Offset of each record batch.
     pub batch_offsets: Vec<i32>,
@@ -197,7 +198,7 @@ impl Metadata {
 }
 
 /// Metadata about the statistics
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, DeepSizeOf)]
 pub struct StatisticsMetadata {
     /// Schema of the page-level statistics.
     ///

rust/lance-file/src/page_table.rs

@@ -4,6 +4,7 @@
 use arrow_array::builder::Int64Builder;
 use arrow_array::{Array, Int64Array};
 use arrow_schema::DataType;
+use deepsize::DeepSizeOf;
 use lance_io::encodings::plain::PlainDecoder;
 use lance_io::encodings::Decoder;
 use snafu::{location, Location};
@@ -13,7 +14,7 @@ use tokio::io::AsyncWriteExt;
 use lance_core::{Error, Result};
 use lance_io::traits::{Reader, Writer};
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, DeepSizeOf)]
 pub struct PageInfo {
     pub position: usize,
     pub length: usize,
@@ -27,7 +28,7 @@ impl PageInfo {
 
 /// Page lookup table.
 ///
-#[derive(Debug, Default, Clone, PartialEq)]
+#[derive(Debug, Default, Clone, PartialEq, DeepSizeOf)]
 pub struct PageTable {
     /// map[field-id,  map[batch-id, PageInfo]]
     pages: BTreeMap<i32, BTreeMap<i32, PageInfo>>,