feat(query): make lazy row fetch page-aware

src/query/service/tests/it/storages/fuse/operations/prewhere.rs

@@ -329,6 +329,7 @@ async fn prepare_prewhere_data() -> Result<PrewhereTestSetup> {
     let part = FuseBlockPartInfo {
         location: "test_block".to_string(),
         bloom_filter_index_location: None,
+        file_size: parquet_bytes.len() as u64,
         bloom_filter_index_size: 0,
         create_on: None,
         nums_rows: num_rows,

src/query/storages/common/cache/Cargo.toml

@@ -14,6 +14,7 @@ databend-common-cache = { workspace = true }
 databend-common-catalog = { workspace = true }
 databend-common-config = { workspace = true }
 databend-common-exception = { workspace = true }
+databend-common-expression = { workspace = true }
 databend-common-metrics = { workspace = true }
 databend-storages-common-index = { workspace = true }
 databend-storages-common-table-meta = { workspace = true }

src/query/storages/common/cache/src/caches.rs

@@ -18,6 +18,7 @@ use std::time::Instant;
 
 use arrow::array::ArrayRef;
 use databend_common_cache::MemSized;
+use databend_common_expression::DataBlock;
 
 use crate::CacheAccessor;
 use crate::InMemoryLruCache;
@@ -69,10 +70,53 @@ pub type PrunePartitionsCache = InMemoryLruCache<(PartStatistics, Partitions)>;
 
 pub type IcebergTableCache = InMemoryLruCache<(Arc<dyn Table>, AtomicBool, Instant)>;
 
-/// In memory object cache of table column array
-pub type ColumnArrayCache = InMemoryLruCache<SizedColumnArray>;
-pub type ArrayRawDataUncompressedSize = usize;
-pub type SizedColumnArray = (ArrayRef, ArrayRawDataUncompressedSize);
+/// In-memory cache of decoded table data.
+pub type TableDataCache = InMemoryLruCache<TableDataCacheEntry>;
+
+pub enum TableDataCacheValue {
+    ColumnArray(ArrayRef),
+    DataBlock(DataBlock),
+}
+
+pub struct TableDataCacheEntry {
+    value: TableDataCacheValue,
+    memory_size: usize,
+}
+
+impl TableDataCacheEntry {
+    pub fn from_column_array(value: ArrayRef, memory_size: usize) -> Self {
+        Self {
+            value: TableDataCacheValue::ColumnArray(value),
+            memory_size,
+        }
+    }
+
+    pub fn from_data_block(value: DataBlock) -> Self {
+        let memory_size = value.memory_size();
+        Self {
+            value: TableDataCacheValue::DataBlock(value),
+            memory_size,
+        }
+    }
+
+    pub fn as_column_array(&self) -> Option<&ArrayRef> {
+        match &self.value {
+            TableDataCacheValue::ColumnArray(value) => Some(value),
+            TableDataCacheValue::DataBlock(_) => None,
+        }
+    }
+
+    pub fn as_data_block(&self) -> Option<&DataBlock> {
+        match &self.value {
+            TableDataCacheValue::ColumnArray(_) => None,
+            TableDataCacheValue::DataBlock(value) => Some(value),
+        }
+    }
+
+    pub fn memory_size(&self) -> usize {
+        self.memory_size
+    }
+}
 
 // Bind Type of cached objects to Caches
 //
@@ -359,10 +403,10 @@ impl From<(PartStatistics, Partitions)> for CacheValue<(PartStatistics, Partitio
     }
 }
 
-impl From<SizedColumnArray> for CacheValue<SizedColumnArray> {
-    fn from(value: SizedColumnArray) -> Self {
+impl From<TableDataCacheEntry> for CacheValue<TableDataCacheEntry> {
+    fn from(value: TableDataCacheEntry) -> Self {
         CacheValue {
-            mem_bytes: value.1,
+            mem_bytes: value.memory_size,
             inner: Arc::new(value),
         }
     }
@@ -384,3 +428,35 @@ impl<T> MemSized for CacheValue<T> {
         self.mem_bytes
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use arrow::array::Int32Array;
+
+    use super::*;
+    use crate::CacheAccessor;
+
+    #[test]
+    fn test_table_data_cache_mixed_entries() {
+        let cache = TableDataCache::with_bytes_capacity("test_table_data".to_string(), 1024);
+
+        let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
+        cache.insert(
+            "column".to_string(),
+            TableDataCacheEntry::from_column_array(array.clone(), 128),
+        );
+        let cached_array = cache.get("column").unwrap();
+        assert!(cached_array.as_data_block().is_none());
+        assert_eq!(cached_array.as_column_array().unwrap().len(), array.len());
+        assert_eq!(cached_array.memory_size(), 128);
+
+        cache.insert(
+            "block".to_string(),
+            TableDataCacheEntry::from_data_block(DataBlock::empty_with_rows(3)),
+        );
+        let cached_block = cache.get("block").unwrap();
+        assert!(cached_block.as_column_array().is_none());
+        assert_eq!(cached_block.as_data_block().unwrap().num_rows(), 3);
+        assert_eq!(cached_block.memory_size(), 0);
+    }
+}

src/query/storages/common/cache/src/manager.rs

@@ -36,7 +36,6 @@ use crate::caches::BlockMetaCache;
 use crate::caches::BloomIndexFilterCache;
 use crate::caches::BloomIndexMetaCache;
 use crate::caches::CacheValue;
-use crate::caches::ColumnArrayCache;
 use crate::caches::ColumnDataCache;
 use crate::caches::ColumnOrientedSegmentInfoCache;
 use crate::caches::CompactSegmentInfoCache;
@@ -48,6 +47,7 @@ use crate::caches::PrunePartitionsCache;
 use crate::caches::SegmentBlockMetasCache;
 use crate::caches::SpatialIndexFileCache;
 use crate::caches::SpatialIndexMetaCache;
+use crate::caches::TableDataCache;
 use crate::caches::TableSnapshotCache;
 use crate::caches::TableSnapshotStatisticCache;
 use crate::caches::VectorIndexFileCache;
@@ -118,7 +118,7 @@ pub struct CacheManager {
     virtual_column_meta_cache: CacheSlot<VirtualColumnMetaCache>,
     prune_partitions_cache: CacheSlot<PrunePartitionsCache>,
     parquet_meta_data_cache: CacheSlot<ParquetMetaDataCache>,
-    in_memory_table_data_cache: CacheSlot<ColumnArrayCache>,
+    in_memory_table_data_cache: CacheSlot<TableDataCache>,
     segment_block_metas_cache: CacheSlot<SegmentBlockMetasCache>,
     block_meta_cache: CacheSlot<BlockMetaCache>,
 
@@ -835,7 +835,7 @@ impl CacheManager {
         self.get_hybrid_cache(self.column_data_cache.get())
     }
 
-    pub fn get_table_data_array_cache(&self) -> Option<ColumnArrayCache> {
+    pub fn get_table_data_cache(&self) -> Option<TableDataCache> {
         self.in_memory_table_data_cache.get()
     }
 
@@ -1023,6 +1023,8 @@ mod tests {
 
     use super::*;
     use crate::ColumnData;
+    use crate::TableDataCacheEntry;
+
     fn config_with_disk_cache_enabled(cache_path: &str) -> CacheConfig {
         CacheConfig {
             data_cache_storage: CacheStorageTypeInnerConfig::Disk,
@@ -1259,9 +1261,12 @@ mod tests {
         // ----- POPULATE BASIC CACHES -----
 
         // 1. Populate in-memory table data cache
-        let in_memory_table_data_cache = cache_manager.get_table_data_array_cache().unwrap();
+        let in_memory_table_data_cache = cache_manager.get_table_data_cache().unwrap();
         let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
-        in_memory_table_data_cache.insert("not matter".to_string(), (array, 1));
+        in_memory_table_data_cache.insert(
+            "not matter".to_string(),
+            TableDataCacheEntry::from_column_array(array, 1),
+        );
 
         // 2. Populate segment block metas cache
         let segment_block_metas_cache = cache_manager.get_segment_block_metas_cache().unwrap();
@@ -1341,7 +1346,7 @@ mod tests {
 
         // ----- VERIFY INITIAL CACHE STATE -----
         // Verify all caches are correctly populated
-        assert!(!cache_manager.get_table_data_array_cache().is_empty());
+        assert!(!cache_manager.get_table_data_cache().is_empty());
         assert!(
             !cache_manager
                 .get_segment_block_metas_cache()

src/query/storages/fuse/Cargo.toml

@@ -25,6 +25,7 @@ databend-common-pipeline-transforms = { workspace = true }
 databend-common-sql = { workspace = true }
 databend-common-statistics = { workspace = true }
 databend-common-storage = { workspace = true }
+databend-common-storages-parquet = { workspace = true }
 databend-common-tracing = { workspace = true }
 databend-common-users = { workspace = true }
 databend-enterprise-fail-safe = { workspace = true }

src/query/storages/fuse/src/fuse_part.rs

@@ -45,6 +45,8 @@ pub struct FuseBlockPartInfo {
 
     pub create_on: Option<DateTime<Utc>>,
     pub nums_rows: usize,
+    #[serde(default)]
+    pub file_size: u64,
     pub columns_meta: HashMap<ColumnId, ColumnMeta>,
     pub columns_stat: Option<HashMap<ColumnId, ColumnStatistics>>,
     pub compression: Compression,
@@ -96,6 +98,7 @@ impl FuseBlockPartInfo {
             bloom_filter_index_size,
             create_on,
             columns_meta,
+            file_size: 0,
             nums_rows: rows_count as usize,
             compression,
             sort_min_max,

src/query/storages/fuse/src/io/read/block/block_reader.rs

@@ -272,7 +272,7 @@ impl BlockReadContext {
         let read_from_in_mem_cache_array: usize = block_read_res
             .cached_column_array
             .iter()
-            .map(|(_, sized_array)| sized_array.1)
+            .map(|(_, cache_entry)| cache_entry.memory_size())
             .sum();
 
         cache_metrics.add_cache_metrics(

src/query/storages/fuse/src/io/read/block/block_reader_merge_io.rs

@@ -19,18 +19,18 @@ use bytes::Bytes;
 use databend_common_exception::Result;
 use databend_common_expression::ColumnId;
 use databend_storages_common_cache::ColumnData;
-use databend_storages_common_cache::SizedColumnArray;
+use databend_storages_common_cache::TableDataCacheEntry;
 use databend_storages_common_io::MergeIOReadResult;
 use enum_as_inner::EnumAsInner;
 use opendal::Buffer;
 
 type CachedColumnData = Vec<(ColumnId, Arc<ColumnData>)>;
-type CachedColumnArray = Vec<(ColumnId, Arc<SizedColumnArray>)>;
+type CachedColumnArray = Vec<(ColumnId, Arc<TableDataCacheEntry>)>;
 
 #[derive(EnumAsInner, Clone)]
 pub enum DataItem<'a> {
     RawData(Buffer),
-    ColumnArray(&'a Arc<SizedColumnArray>),
+    ColumnArray(&'a Arc<TableDataCacheEntry>),
 }
 
 pub struct BlockReadResult {

src/query/storages/fuse/src/io/read/block/block_reader_merge_io_async.rs

@@ -64,7 +64,7 @@ impl BlockReadContext {
         let mut ranges = vec![];
         // for async read, try using table data cache (if enabled in settings)
         let column_data_cache = CacheManager::instance().get_column_data_cache();
-        let column_array_cache = CacheManager::instance().get_table_data_array_cache();
+        let table_data_cache = CacheManager::instance().get_table_data_cache();
         let mut cached_column_data = vec![];
         let mut cached_column_array = vec![];
 
@@ -84,14 +84,16 @@ impl BlockReadContext {
 
                 // first, check in memory table data cache
                 // column_array_cache
-                if let Some(cache_array) = column_array_cache.get_sized(&column_cache_key, len) {
-                    // Record bytes scanned from memory cache (table data only)
-                    Profile::record_usize_profile(
-                        ProfileStatisticsName::ScanBytesFromMemory,
-                        len as usize,
-                    );
-                    cached_column_array.push((*column_id, cache_array));
-                    continue;
+                if let Some(cache_entry) = table_data_cache.get_sized(&column_cache_key, len) {
+                    if cache_entry.as_column_array().is_some() {
+                        // Record bytes scanned from memory cache (table data only)
+                        Profile::record_usize_profile(
+                            ProfileStatisticsName::ScanBytesFromMemory,
+                            len as usize,
+                        );
+                        cached_column_array.push((*column_id, cache_entry));
+                        continue;
+                    }
                 }
 
                 // and then, check on disk table data cache