Skip to content

fix: Validate file cache identity and isolate object stores #25858

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
limadog9 wants to merge 1 commit into
base: main
Choose a base branch
from
+1,591 −286
Open

fix: Validate file cache identity and isolate object stores #25858

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions datafusion-cli/src/functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -701,6 +701,7 @@ impl TableFunctionImpl for MetadataCacheFunc {

let schema = Arc::new(Schema::new(vec![
Field::new("path", DataType::Utf8, false),
Field::new("object_store_url", DataType::Utf8, false),
Field::new(
"file_modified",
DataType::Timestamp(TimeUnit::Millisecond, None),
Expand All @@ -718,6 +719,7 @@ impl TableFunctionImpl for MetadataCacheFunc {

// construct record batch from metadata
let mut path_arr = Vec::with_capacity(cached_entries.len());
let mut object_store_url_arr = Vec::with_capacity(cached_entries.len());
let mut file_modified_arr = Vec::with_capacity(cached_entries.len());
let mut file_size_bytes_arr = Vec::with_capacity(cached_entries.len());
let mut e_tag_arr = Vec::with_capacity(cached_entries.len());
Expand All @@ -728,7 +730,8 @@ impl TableFunctionImpl for MetadataCacheFunc {
let mut extra = Vec::new();

for (path, entry) in cached_entries {
path_arr.push(path.to_string());
path_arr.push(path.path.to_string());
object_store_url_arr.push(path.object_store_url.to_string());
file_modified_arr
.push(Some(entry.value.meta.last_modified.timestamp_millis()));
file_size_bytes_arr.push(entry.value.meta.size);
Expand All @@ -754,6 +757,7 @@ impl TableFunctionImpl for MetadataCacheFunc {
schema.clone(),
vec![
Arc::new(StringArray::from(path_arr)),
Arc::new(StringArray::from(object_store_url_arr)),
Arc::new(TimestampMillisecondArray::from(file_modified_arr)),
Arc::new(UInt64Array::from(file_size_bytes_arr)),
Arc::new(StringArray::from(e_tag_arr)),
Expand Down Expand Up @@ -821,6 +825,7 @@ impl TableFunctionImpl for StatisticsCacheFunc {

let schema = Arc::new(Schema::new(vec![
Field::new("path", DataType::Utf8, false),
Field::new("object_store_url", DataType::Utf8, false),
Field::new("table", DataType::Utf8, false),
Field::new(
"file_modified",
Expand All @@ -839,6 +844,7 @@ impl TableFunctionImpl for StatisticsCacheFunc {

// construct record batch from metadata
let mut path_arr = vec![];
let mut object_store_url_arr = vec![];
let mut table_arr = vec![];
let mut file_modified_arr = vec![];
let mut file_size_bytes_arr = vec![];
Expand All @@ -853,7 +859,8 @@ impl TableFunctionImpl for StatisticsCacheFunc {
if let Some(file_statistics_cache) = self.cache_manager.get_file_statistic_cache()
{
for (path, entry) in file_statistics_cache.list_entries() {
path_arr.push(path.path.to_string());
path_arr.push(path.store_path.path.to_string());
object_store_url_arr.push(path.store_path.object_store_url.to_string());
table_arr
.push(path.table.map_or_else(|| "".to_string(), |t| t.to_string()));
file_modified_arr
Expand Down Expand Up @@ -881,6 +888,7 @@ impl TableFunctionImpl for StatisticsCacheFunc {
schema.clone(),
vec![
Arc::new(StringArray::from(path_arr)),
Arc::new(StringArray::from(object_store_url_arr)),
Arc::new(StringArray::from(table_arr)),
Arc::new(TimestampMillisecondArray::from(file_modified_arr)),
Arc::new(UInt64Array::from(file_size_bytes_arr)),
Expand Down Expand Up @@ -912,11 +920,13 @@ impl TableFunctionImpl for StatisticsCacheFunc {
/// +---------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------+
/// | column_name | data_type | is_nullable |
/// +---------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------+
/// | table | Utf8 | NO |
/// | table | Utf8 | YES |
/// | path | Utf8 | NO |
/// | object_store_url | Utf8 | NO |
/// | metadata_size_bytes | UInt64 | NO |
/// | expires_in | Duration(ms) | YES |
/// | metadata_list | List(Struct("file_path": non-null Utf8, "file_modified": non-null Timestamp(ms), "file_size_bytes": non-null UInt64, "e_tag": Utf8, "version": Utf8), field: 'metadata') | YES |
/// | hits | UInt64 | NO |
/// +---------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------+
/// ```
#[derive(Debug)]
Expand Down Expand Up @@ -986,6 +996,7 @@ impl TableFunctionImpl for ListFilesCacheFunc {
let schema = Arc::new(Schema::new(vec![
Field::new("table", DataType::Utf8, true),
Field::new("path", DataType::Utf8, false),
Field::new("object_store_url", DataType::Utf8, false),
Field::new("metadata_size_bytes", DataType::UInt64, false),
// expires field in ListFilesEntry has type Instant when set, from which we cannot get "the number of seconds", hence using Duration instead of Timestamp as data type.
Field::new(
Expand All @@ -1003,6 +1014,7 @@ impl TableFunctionImpl for ListFilesCacheFunc {

let mut table_arr = vec![];
let mut path_arr = vec![];
let mut object_store_url_arr = vec![];
let mut metadata_size_bytes_arr = vec![];
let mut expires_arr = vec![];

Expand All @@ -1020,7 +1032,8 @@ impl TableFunctionImpl for ListFilesCacheFunc {

for (path, entry) in list_files_cache.list_entries() {
table_arr.push(path.table.map(|t| t.to_string()));
path_arr.push(path.path.to_string());
path_arr.push(path.store_path.path.to_string());
object_store_url_arr.push(path.store_path.object_store_url.to_string());
metadata_size_bytes_arr.push(entry.size_bytes as u64);
// calculates time left before entry expires
expires_arr.push(
Expand Down Expand Up @@ -1062,6 +1075,7 @@ impl TableFunctionImpl for ListFilesCacheFunc {
vec![
Arc::new(StringArray::from(table_arr)),
Arc::new(StringArray::from(path_arr)),
Arc::new(StringArray::from(object_store_url_arr)),
Arc::new(UInt64Array::from(metadata_size_bytes_arr)),
Arc::new(DurationMillisecondArray::from(expires_arr)),
Arc::new(GenericListArray::new(
Expand Down
75 changes: 75 additions & 0 deletions datafusion-cli/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -499,6 +499,8 @@ mod tests {
use std::time::Duration;

use super::*;
use arrow::array::{ArrayRef, RecordBatch, StringArray};
use arrow::datatypes::DataType;
use datafusion::execution::cache::default_cache::DefaultCache;
use datafusion::{
common::test_util::batches_to_string,
Expand Down Expand Up @@ -685,6 +687,79 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn test_cache_functions_object_store_identity() -> Result<(), DataFusionError> {
let ctx = SessionContext::new();
let cache_manager = &ctx.runtime_env().cache_manager;
ctx.register_udtf(
"metadata_cache",
Arc::new(MetadataCacheFunc::new(Arc::clone(cache_manager))),
);
ctx.register_udtf(
"statistics_cache",
Arc::new(StatisticsCacheFunc::new(Arc::clone(cache_manager))),
);
ctx.register_udtf(
"list_files_cache",
Arc::new(ListFilesCacheFunc::new(Arc::clone(cache_manager))),
);

for store in ["first", "second"] {
ctx.register_object_store(
&Url::parse(&format!("mem://{store}")).unwrap(),
Arc::new(InMemory::new()),
);
ctx.sql(&format!(
"COPY (SELECT 1 AS value) TO 'mem://{store}/data/file.parquet' STORED AS PARQUET"
))
.await?
.collect()
.await?;
ctx.read_parquet(
format!("mem://{store}/data/"),
ParquetReadOptions::default(),
)
.await?
.collect()
.await?;
}

for (function, expected_path) in [
("metadata_cache", "data/file.parquet"),
("statistics_cache", "data/file.parquet"),
("list_files_cache", "data"),
] {
let df = ctx
.sql(&format!(
"SELECT object_store_url, path FROM {function}() ORDER BY object_store_url"
))
.await?;
let field = df
.schema()
.field_with_unqualified_name("object_store_url")?;
assert_eq!(field.data_type(), &DataType::Utf8, "{function}");
assert!(!field.is_nullable(), "{function}");
let results = df.collect().await?;
let expected = RecordBatch::try_from_iter([
(
"object_store_url",
Arc::new(StringArray::from(vec!["mem://first/", "mem://second/"]))
as ArrayRef,
),
(
"path",
Arc::new(StringArray::from(vec![expected_path, expected_path])),
),
])?;
assert_eq!(
batches_to_string(&results),
batches_to_string(&[expected]),
"{function}"
);
}
Ok(())
}

#[tokio::test]
async fn test_metadata_cache() -> Result<(), DataFusionError> {
let ctx = SessionContext::new();
Expand Down
8 changes: 5 additions & 3 deletions datafusion-examples/examples/custom_data_source/custom_file_format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ use datafusion::{
table_schema::TableSchema,
},
error::Result,
execution::session_state::SessionStateBuilder,
execution::{object_store::ObjectStoreUrl, session_state::SessionStateBuilder},
physical_expr_common::sort_expr::LexRequirement,
physical_plan::ExecutionPlan,
prelude::SessionContext,
Expand Down Expand Up @@ -124,22 +124,24 @@ impl FileFormat for TSVFileFormat {
&self,
state: &dyn Session,
store: &Arc<dyn ObjectStore>,
object_store_url: &ObjectStoreUrl,
objects: &[ObjectMeta],
) -> Result<SchemaRef> {
self.csv_file_format
.infer_schema(state, store, objects)
.infer_schema(state, store, object_store_url, objects)
.await
}

async fn infer_stats(
&self,
state: &dyn Session,
store: &Arc<dyn ObjectStore>,
object_store_url: &ObjectStoreUrl,
table_schema: SchemaRef,
object: &ObjectMeta,
) -> Result<Statistics> {
self.csv_file_format
.infer_stats(state, store, table_schema, object)
.infer_stats(state, store, object_store_url, table_schema, object)
.await
}

Expand Down
5 changes: 4 additions & 1 deletion datafusion/catalog-listing/src/options.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,7 +300,10 @@ impl ListingOptions {
.filter(|object_meta| object_meta.size > 0)
.collect();

let schema = self.format.infer_schema(state, &store, &files).await?;
let schema = self
.format
.infer_schema(state, &store, &table_path.object_store(), &files)
.await?;

Ok(schema)
}
Expand Down
51 changes: 36 additions & 15 deletions datafusion/catalog-listing/src/table.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,10 @@ use datafusion_datasource::{
ListingTableUrl, PartitionedFile, TableSchemaBuilder, compute_all_files_statistics,
};
use datafusion_execution::cache::cache_manager::{
CachedFileMetadata, FileStatisticsCache, SchemaFingerprint, TableScopedPath,
CachedFileMetadata, FileStatisticsCache, SchemaFingerprint, StoreScopedPath,
TableScopedPath,
};
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_expr::dml::InsertOp;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
Expand Down Expand Up @@ -867,13 +869,15 @@ impl ListingTable {
lfc.drop_table_entries(table_ref)?;
} else {
let table_prefix = table_path.prefix();
let object_store_url = table_path.object_store();
let keys: Vec<_> = lfc
.list_entries()
.into_keys()
.filter(|key| {
key.table.is_none()
&& (key.path.prefix_matches(table_prefix)
|| table_prefix.prefix_matches(&key.path))
&& key.store_path.object_store_url == object_store_url
&& (key.store_path.path.prefix_matches(table_prefix)
|| table_prefix.prefix_matches(&key.store_path.path))
})
.collect();
for key in keys {
Expand Down Expand Up @@ -937,6 +941,7 @@ impl ListingTable {
&'a self,
ctx: &'a dyn Session,
store: &'a Arc<dyn ObjectStore>,
object_store_url: &'a ObjectStoreUrl,
listing_time_filters: &'a [Expr],
file_limit: Option<usize>,
) -> datafusion_common::Result<(FileGroup, bool)> {
Expand Down Expand Up @@ -968,8 +973,13 @@ impl ListingTable {
.map(|part_file| async {
let part_file = part_file?;
let (statistics, ordering) = if ctx.config().collect_statistics() {
self.do_collect_statistics_and_ordering(ctx, store, &part_file)
.await?
self.do_collect_statistics_and_ordering(
ctx,
store,
object_store_url,
&part_file,
)
.await?
} else {
(Arc::new(Statistics::new_unknown(&self.file_schema)), None)
};
Expand Down Expand Up @@ -998,8 +1008,8 @@ impl ListingTable {
);
}

let store = if let Some(url) = self.table_paths.first() {
ctx.runtime_env().object_store(url)?
let (store, object_store_url) = if let Some(url) = self.table_paths.first() {
(ctx.runtime_env().object_store(url)?, url.object_store())
} else {
return Ok(ListFilesResult {
file_groups: vec![],
Expand All @@ -1008,7 +1018,7 @@ impl ListingTable {
});
};
let (file_group, inexact_stats) = self
.collect_files_for_scan(ctx, &store, filters, limit)
.collect_files_for_scan(ctx, &store, &object_store_url, filters, limit)
.await?;

// Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N
Expand Down Expand Up @@ -1060,17 +1070,18 @@ impl ListingTable {
);
}

let store = if let Some(url) = self.table_paths.first() {
ctx.runtime_env().object_store(url)?
let (store, object_store_url) = if let Some(url) = self.table_paths.first() {
(ctx.runtime_env().object_store(url)?, url.object_store())
} else {
return Ok(ListFilesResult {
file_groups: vec![],
statistics: Statistics::new_unknown(&self.file_schema),
grouped_by_partition: false,
});
};
let (file_group, inexact_stats) =
self.collect_files_for_scan(ctx, &store, &[], None).await?;
let (file_group, inexact_stats) = self
.collect_files_for_scan(ctx, &store, &object_store_url, &[], None)
.await?;
let mut file_groups = file_group.split_files(file_group_count);
if !file_groups.is_empty() {
file_groups.resize_with(file_group_count, || FileGroup::new(vec![]));
Expand Down Expand Up @@ -1141,15 +1152,19 @@ impl ListingTable {
&self,
ctx: &dyn Session,
store: &Arc<dyn ObjectStore>,
object_store_url: &ObjectStoreUrl,
part_file: &PartitionedFile,
) -> datafusion_common::Result<(Arc<Statistics>, Option<LexOrdering>)> {
let path = TableScopedPath {
table: part_file.table_reference.clone(),
path: part_file.object_meta.location.clone(),
store_path: StoreScopedPath::new(
object_store_url.clone(),
part_file.object_meta.location.clone(),
),
};
let meta = &part_file.object_meta;

// Check cache first. The key stays `{table, path}` for cheap lookups;
// Check cache first. The key identifies the store, table and path;
// the cached value carries the schema fingerprint to prevent reusing
// stats computed under a different file schema.
if let Some(cache) = &self.collected_statistics
Expand All @@ -1164,7 +1179,13 @@ impl ListingTable {
let file_meta = self
.options
.format
.infer_stats_and_ordering(ctx, store, Arc::clone(&self.file_schema), meta)
.infer_stats_and_ordering(
ctx,
store,
object_store_url,
Arc::clone(&self.file_schema),
meta,
)
.await?;

let statistics = Arc::new(file_meta.statistics);
Expand Down
Loading
Loading