Skip to content

Commit 044cc2b

Browse files
HarishNarasimhanKHarish Narasimhan
authored and
bkhishor
committed
Fix: use NativeCacheManagerHandle instead of raw cache manager pointer (opensearch-project#21577)
Signed-off-by: Harish Narasimhan <hxarishk@amazon.com> Co-authored-by: Harish Narasimhan <hxarishk@amazon.com>
1 parent 2fe8f0d commit 044cc2b

15 files changed

Lines changed: 283 additions & 126 deletions

File tree

‎sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -738,6 +738,9 @@ pub unsafe fn sql_to_substrait(
738738
.with_file_metadata_cache(Some(
739739
runtime.runtime_env.cache_manager.get_file_metadata_cache(),
740740
))
741+
.with_metadata_cache_limit(
742+
runtime.runtime_env.cache_manager.get_metadata_cache_limit(),
743+
)
741744
.with_files_statistics_cache(
742745
runtime.runtime_env.cache_manager.get_file_statistic_cache(),
743746
),

‎sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ use datafusion::execution::cache::cache_manager::{FileMetadataCache, FileStatist
1111
use datafusion::execution::cache::cache_unit::DefaultFileStatisticsCache;
1212
use datafusion::execution::cache::CacheAccessor;
1313
use crate::statistics_cache::compute_parquet_statistics;
14-
use tokio::runtime::Runtime;
1514
use crate::cache::MutexFileMetadataCache;
1615
use crate::statistics_cache::CustomStatisticsCache;
1716
use object_store::path::Path;
@@ -106,15 +105,15 @@ impl CustomCacheManager {
106105
}
107106

108107
/// Add multiple files to all applicable caches
109-
pub fn add_files(&self, file_paths: &[String]) -> Result<Vec<(String, bool)>, String> {
108+
pub fn add_files(&self, file_paths: &[String], rt_handle: &tokio::runtime::Handle) -> Result<Vec<(String, bool)>, String> {
110109
let mut results = Vec::new();
111110

112111
for file_path in file_paths {
113112
let mut any_success = false;
114113
let mut errors = Vec::new();
115114

116115
// Add to metadata cache
117-
match self.metadata_cache_put(file_path) {
116+
match self.metadata_cache_put(file_path, rt_handle) {
118117
Ok(true) => {
119118
any_success = true;
120119
}
@@ -185,7 +184,6 @@ impl CustomCacheManager {
185184
// Remove from statistics cache
186185
if let Some(cache) = &self.statistics_cache {
187186
let path = Path::from(file_path.clone());
188-
// Use the CacheAccessor remove method to properly update memory tracking
189187
if cache.remove(&path).is_some() {
190188
any_removed = true;
191189
}
@@ -342,7 +340,7 @@ impl CustomCacheManager {
342340
}
343341

344342
/// Internal method to put metadata into cache
345-
fn metadata_cache_put(&self, file_path: &str) -> Result<bool, String> {
343+
fn metadata_cache_put(&self, file_path: &str, rt_handle: &tokio::runtime::Handle) -> Result<bool, String> {
346344
if !file_path.to_lowercase().ends_with(".parquet") {
347345
return Ok(false); // Skip unsupported formats
348346
}
@@ -367,16 +365,14 @@ impl CustomCacheManager {
367365
// 2. Load the complete metadata including column and offset indexes
368366
// 3. Automatically put the metadata into the cache (lines 155-160 in datafusion's metadata.rs)
369367
// This ensures we cache exactly what DataFusion would cache during query execution
370-
let _parquet_metadata = Runtime::new()
371-
.map_err(|e| format!("Failed to create Tokio Runtime: {}", e))?
372-
.block_on(async {
373-
let df_metadata = DFParquetMetadata::new(store.as_ref(), object_meta)
374-
.with_file_metadata_cache(Some(metadata_cache));
375-
376-
// fetch_metadata() performs the cache put operation internally
377-
df_metadata.fetch_metadata().await
378-
.map_err(|e| format!("Failed to fetch metadata: {}", e))
379-
})?;
368+
let _parquet_metadata = rt_handle.block_on(async {
369+
let df_metadata = DFParquetMetadata::new(store.as_ref(), object_meta)
370+
.with_file_metadata_cache(Some(metadata_cache));
371+
372+
// fetch_metadata() performs the cache put operation internally
373+
df_metadata.fetch_metadata().await
374+
.map_err(|e| format!("Failed to fetch metadata: {}", e))
375+
})?;
380376

381377
// Verify the metadata was cached properly
382378
match cache_ref.inner.lock() {

‎sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -547,8 +547,11 @@ pub unsafe extern "C" fn df_cache_manager_add_files(
547547
);
548548
}
549549

550-
manager
551-
.add_files(&file_paths)
550+
let rt_manager = get_rt_manager()
551+
.map_err(|e| format!("df_cache_manager_add_files: {}", e))?;
552+
let rt_handle = rt_manager.io_runtime.handle();
553+
554+
manager.add_files(&file_paths, rt_handle)
552555
.map_err(|e| format!("df_cache_manager_add_files: {}", e))?;
553556
Ok(0)
554557
}

‎sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,9 @@ pub async fn execute_indexed_query(
114114
.with_file_metadata_cache(Some(
115115
runtime.runtime_env.cache_manager.get_file_metadata_cache(),
116116
))
117+
.with_metadata_cache_limit(
118+
runtime.runtime_env.cache_manager.get_metadata_cache_limit(),
119+
)
117120
.with_files_statistics_cache(
118121
runtime.runtime_env.cache_manager.get_file_statistic_cache(),
119122
),
@@ -449,9 +452,17 @@ pub async unsafe fn execute_indexed_with_context(
449452
let state = ctx.state();
450453
let store = state.runtime_env().object_store(&table_path)?;
451454

452-
let (segments, schema) = build_segments(&state, Arc::clone(&store), object_metas.as_ref(), writer_generations.as_ref())
453-
.await
454-
.map_err(DataFusionError::Execution)?;
455+
let metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache();
456+
457+
let (segments, schema) = build_segments(
458+
&state,
459+
Arc::clone(&store),
460+
object_metas.as_ref(),
461+
writer_generations.as_ref(),
462+
metadata_cache,
463+
)
464+
.await
465+
.map_err(DataFusionError::Execution)?;
455466
let schema = crate::schema_coerce::coerce_inferred_schema(schema);
456467
for (i, seg) in segments.iter().enumerate() {
457468
}

‎sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs‎

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,17 @@ use std::sync::Arc;
2222

2323
use datafusion::arrow::datatypes::SchemaRef;
2424
use datafusion::common::Result;
25+
use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata;
2526
use datafusion::datasource::physical_plan::parquet::{
2627
ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory, RowGroupAccess,
2728
};
2829
use datafusion::datasource::physical_plan::ParquetSource;
30+
use datafusion::execution::cache::cache_manager::FileMetadataCache;
2931
use datafusion::execution::object_store::ObjectStoreUrl;
3032
use datafusion::execution::SendableRecordBatchStream;
31-
use datafusion::parquet::arrow::arrow_reader::{
32-
ArrowReaderMetadata, ArrowReaderOptions, RowSelection,
33-
};
33+
use datafusion::parquet::arrow::arrow_reader::{ArrowReaderOptions, RowSelection};
3434
use datafusion::parquet::arrow::async_reader::AsyncFileReader;
35+
use datafusion::parquet::arrow::parquet_to_arrow_schema;
3536
use datafusion::parquet::file::metadata::ParquetMetaData;
3637
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
3738
use datafusion::physical_plan::ExecutionPlan;
@@ -45,28 +46,30 @@ use prost::bytes::Bytes;
4546

4647
// ── Parquet Metadata Loading ─────────────────────────────────────────
4748

48-
/// Load parquet metadata with page index over the object store.
49+
/// Load parquet metadata via DataFusion's `DFParquetMetadata`, consulting the
50+
/// caller-supplied `FileMetadataCache`.
4951
pub async fn load_parquet_metadata(
5052
store: Arc<dyn ObjectStore>,
5153
location: &object_store::path::Path,
54+
metadata_cache: Arc<dyn FileMetadataCache>,
5255
) -> std::result::Result<(SchemaRef, u64, Arc<ParquetMetaData>), String> {
5356
let meta = store
5457
.head(location)
5558
.await
5659
.map_err(|e| format!("object-store head {}: {}", location, e))?;
5760
let size = meta.size;
58-
let mut reader =
59-
datafusion::parquet::arrow::async_reader::ParquetObjectReader::new(store, location.clone())
60-
.with_file_size(size);
61-
let options = ArrowReaderOptions::new().with_page_index(true);
62-
let arrow_metadata = ArrowReaderMetadata::load_async(&mut reader, options)
61+
62+
let pq_meta = DFParquetMetadata::new(&*store, &meta)
63+
.with_file_metadata_cache(Some(metadata_cache))
64+
.fetch_metadata()
6365
.await
6466
.map_err(|e| format!("load parquet metadata {}: {}", location, e))?;
65-
Ok((
66-
arrow_metadata.schema().clone(),
67-
size,
68-
arrow_metadata.metadata().clone(),
69-
))
67+
68+
let file_meta = pq_meta.file_metadata();
69+
let schema = parquet_to_arrow_schema(file_meta.schema_descr(), file_meta.key_value_metadata())
70+
.map_err(|e| format!("parquet_to_arrow_schema {}: {}", location, e))?;
71+
72+
Ok((Arc::new(schema), size, pq_meta))
7073
}
7174

7275
/// Configuration for creating a per-row-group parquet stream.

‎sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/segment_info.rs‎

Lines changed: 60 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use std::sync::Arc;
2727
use datafusion::catalog::Session;
2828
use datafusion::datasource::file_format::parquet::ParquetFormat;
2929
use datafusion::datasource::file_format::FileFormat;
30+
use datafusion::execution::cache::cache_manager::FileMetadataCache;
3031

3132
/// Parquet footer kv key under which the writer stamps the writer generation.
3233
/// Must match `parquet-data-format`'s `WRITER_GENERATION_KEY`.
@@ -50,6 +51,7 @@ pub async fn build_segments(
5051
store: Arc<dyn object_store::ObjectStore>,
5152
object_metas: &[object_store::ObjectMeta],
5253
writer_generations: &[i64],
54+
metadata_cache: Arc<dyn FileMetadataCache>,
5355
) -> Result<(Vec<SegmentFileInfo>, arrow::datatypes::SchemaRef), String> {
5456
if object_metas.len() != writer_generations.len() {
5557
return Err(format!(
@@ -71,10 +73,13 @@ pub async fn build_segments(
7173
// RuntimeEnv that `infer_schema` uses below shares this data when
7274
// both are pointed at the same object location, so this isn't a
7375
// duplicated cold fetch in practice.
74-
let (_file_schema, size, pq_meta) =
75-
parquet_bridge::load_parquet_metadata(Arc::clone(&store), &meta.location)
76-
.await
77-
.map_err(|e| format!("parquet metadata {}: {}", meta.location, e))?;
76+
let (_file_schema, size, pq_meta) = parquet_bridge::load_parquet_metadata(
77+
Arc::clone(&store),
78+
&meta.location,
79+
Arc::clone(&metadata_cache),
80+
)
81+
.await
82+
.map_err(|e| format!("parquet metadata {}: {}", meta.location, e))?;
7883

7984
let mut row_groups = Vec::new();
8085
let mut offset: i64 = 0;
@@ -170,11 +175,18 @@ mod tests {
170175
use super::*;
171176
use arrow::datatypes::{DataType, Field, Schema};
172177
use arrow_array::{Int32Array, RecordBatch, StringArray};
178+
use datafusion::execution::cache::cache_unit::DefaultFilesMetadataCache;
173179
use datafusion::execution::context::SessionContext;
174180
use datafusion::parquet::arrow::ArrowWriter;
175181
use object_store::{local::LocalFileSystem, path::Path as ObjectPath, ObjectStore, ObjectStoreExt};
176182
use tempfile::tempdir;
177183

184+
/// Mirror of what `CacheManager::try_new` auto-installs when no custom
185+
/// metadata cache is configured.
186+
fn default_metadata_cache() -> Arc<dyn FileMetadataCache> {
187+
Arc::new(DefaultFilesMetadataCache::new(50 * 1024 * 1024))
188+
}
189+
178190
/// Write a parquet file with the given schema + arrays into `dir`
179191
/// under `filename`. Returns the absolute filesystem path.
180192
fn write_parquet(
@@ -239,9 +251,15 @@ mod tests {
239251
let metas = object_metas(store.as_ref(), &[p0, p1]).await;
240252
let ctx = SessionContext::new();
241253
let gens: Vec<i64> = (0..metas.len() as i64).collect();
242-
let (segments, merged) = build_segments(&ctx.state(), Arc::clone(&store), &metas, &gens)
243-
.await
244-
.unwrap();
254+
let (segments, merged) = build_segments(
255+
&ctx.state(),
256+
Arc::clone(&store),
257+
&metas,
258+
&gens,
259+
default_metadata_cache(),
260+
)
261+
.await
262+
.unwrap();
245263

246264
assert_eq!(segments.len(), 2);
247265
assert_eq!(merged.fields().len(), 2);
@@ -289,9 +307,15 @@ mod tests {
289307
let metas = object_metas(store.as_ref(), &[p0, p1]).await;
290308
let ctx = SessionContext::new();
291309
let gens: Vec<i64> = (0..metas.len() as i64).collect();
292-
let (_segments, merged) = build_segments(&ctx.state(), Arc::clone(&store), &metas, &gens)
293-
.await
294-
.unwrap();
310+
let (_segments, merged) = build_segments(
311+
&ctx.state(),
312+
Arc::clone(&store),
313+
&metas,
314+
&gens,
315+
default_metadata_cache(),
316+
)
317+
.await
318+
.unwrap();
295319

296320
let names: Vec<&str> = merged.fields().iter().map(|f| f.name().as_str()).collect();
297321
assert_eq!(
@@ -349,12 +373,24 @@ mod tests {
349373
let gens_ab: Vec<i64> = (0..metas_ab.len() as i64).collect();
350374
let gens_ba: Vec<i64> = (0..metas_ba.len() as i64).collect();
351375

352-
let (_, schema_ab) = build_segments(&ctx.state(), Arc::clone(&store), &metas_ab, &gens_ab)
353-
.await
354-
.unwrap();
355-
let (_, schema_ba) = build_segments(&ctx.state(), Arc::clone(&store), &metas_ba, &gens_ba)
356-
.await
357-
.unwrap();
376+
let (_, schema_ab) = build_segments(
377+
&ctx.state(),
378+
Arc::clone(&store),
379+
&metas_ab,
380+
&gens_ab,
381+
default_metadata_cache(),
382+
)
383+
.await
384+
.unwrap();
385+
let (_, schema_ba) = build_segments(
386+
&ctx.state(),
387+
Arc::clone(&store),
388+
&metas_ba,
389+
&gens_ba,
390+
default_metadata_cache(),
391+
)
392+
.await
393+
.unwrap();
358394

359395
let fields_ab: Vec<&str> = schema_ab
360396
.fields()
@@ -398,7 +434,14 @@ mod tests {
398434
let metas = object_metas(store.as_ref(), &[p0, p1]).await;
399435
let ctx = SessionContext::new();
400436
let gens: Vec<i64> = (0..metas.len() as i64).collect();
401-
let result = build_segments(&ctx.state(), Arc::clone(&store), &metas, &gens).await;
437+
let result = build_segments(
438+
&ctx.state(),
439+
Arc::clone(&store),
440+
&metas,
441+
&gens,
442+
default_metadata_cache(),
443+
)
444+
.await;
402445
assert!(
403446
result.is_err(),
404447
"conflicting Int32 / Int64 on same field name must fail at schema-union time"

‎sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,9 @@ pub async fn execute_query(
7171
.with_file_metadata_cache(Some(
7272
runtime.runtime_env.cache_manager.get_file_metadata_cache(),
7373
))
74+
.with_metadata_cache_limit(
75+
runtime.runtime_env.cache_manager.get_metadata_cache_limit(),
76+
)
7477
.with_files_statistics_cache(
7578
runtime.runtime_env.cache_manager.get_file_statistic_cache(),
7679
),
@@ -131,10 +134,16 @@ pub async fn execute_query(
131134
.with_listing_options(listing_options)
132135
.with_schema(resolved_schema);
133136

134-
let provider = Arc::new(ListingTable::try_new(table_config).map_err(|e| {
135-
error!("Failed to create listing table: {}", e);
136-
e
137-
})?);
137+
// Wire the global statistics cache into the ListingTable.
138+
let stats_cache = runtime.runtime_env.cache_manager.get_file_statistic_cache();
139+
let provider = Arc::new(
140+
ListingTable::try_new(table_config)
141+
.map_err(|e| {
142+
error!("Failed to create listing table: {}", e);
143+
e
144+
})?
145+
.with_cache(stats_cache),
146+
);
138147

139148
ctx.register_table(&table_name, provider).map_err(|e| {
140149
error!("Failed to register table: {}", e);

‎sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,9 @@ pub async unsafe fn create_session_context(
9696
.with_file_metadata_cache(Some(
9797
runtime.runtime_env.cache_manager.get_file_metadata_cache(),
9898
))
99+
.with_metadata_cache_limit(
100+
runtime.runtime_env.cache_manager.get_metadata_cache_limit(),
101+
)
99102
.with_files_statistics_cache(
100103
runtime.runtime_env.cache_manager.get_file_statistic_cache(),
101104
),
@@ -172,13 +175,19 @@ pub async unsafe fn create_session_context(
172175
.with_listing_options(listing_options)
173176
.with_schema(resolved_schema);
174177

175-
let provider = Arc::new(ListingTable::try_new(table_config).map_err(|e| {
176-
error!(
177-
"create_session_context: failed to create listing table: {}",
178-
e
179-
);
180-
e
181-
})?);
178+
// Wire the global statistics cache into the ListingTable.
179+
let stats_cache = runtime.runtime_env.cache_manager.get_file_statistic_cache();
180+
let provider = Arc::new(
181+
ListingTable::try_new(table_config)
182+
.map_err(|e| {
183+
error!(
184+
"create_session_context: failed to create listing table: {}",
185+
e
186+
);
187+
e
188+
})?
189+
.with_cache(stats_cache),
190+
);
182191

183192
ctx.register_table(table_name, provider).map_err(|e| {
184193
error!(

0 commit comments

Comments
 (0)