Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
56 changes: 48 additions & 8 deletions crates/integrations/datafusion/src/catalog.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use std::sync::RwLock;
use async_trait::async_trait;
use datafusion::catalog::{CatalogProvider, MemorySchemaProvider, SchemaProvider};
use datafusion::common::{plan_datafusion_err, Column};
use datafusion::datasource::TableProvider;
use datafusion::datasource::{TableProvider, TableType};
use datafusion::error::Result as DFResult;
use datafusion::execution::SessionState;
use datafusion::logical_expr::{expr_fn::cast, Expr, LogicalPlan, LogicalPlanBuilder};
Expand Down Expand Up @@ -338,6 +338,8 @@ pub struct PaimonSchemaProvider {
dynamic_options: DynamicOptions,
/// Optional temporary in-memory provider for temp tables and views.
temp_provider: Option<Arc<MemorySchemaProvider>>,
/// Table types populated together with `table_names` for metadata-only lookups.
catalog_table_types: RwLock<HashMap<String, TableType>>,
blob_reader_registry: BlobReaderRegistry,
session_state: Option<SessionStateProvider>,
schema_force_view_types: bool,
Expand Down Expand Up @@ -369,6 +371,7 @@ impl PaimonSchemaProvider {
database,
dynamic_options,
temp_provider,
catalog_table_types: RwLock::new(HashMap::new()),
blob_reader_registry,
session_state,
schema_force_view_types: true,
Expand All @@ -386,30 +389,46 @@ impl SchemaProvider for PaimonSchemaProvider {
fn table_names(&self) -> Vec<String> {
let catalog = Arc::clone(&self.catalog);
let database = self.database.clone();
let mut names = block_on_with_runtime(
let (mut names, views) = block_on_with_runtime(
{
let db = database.clone();
async move {
let mut names = match catalog.list_tables(&db).await {
let names = match catalog.list_tables(&db).await {
Ok(names) => names,
Err(e) => {
log::error!("failed to list tables in '{}': {e}", db);
vec![]
}
};
match catalog.list_views(&db).await {
Ok(views) => names.extend(views),
Err(paimon::Error::Unsupported { .. }) => {}
let views = match catalog.list_views(&db).await {
Ok(views) => views,
Err(paimon::Error::Unsupported { .. }) => vec![],
Err(error) => {
log::error!("failed to list views in '{}': {error}", db);
vec![]
}
}
names
};
(names, views)
}
},
"paimon catalog access thread panicked",
);

let mut catalog_table_types = HashMap::with_capacity(names.len() + views.len());
for name in &names {
catalog_table_types.insert(name.clone(), TableType::Base);
}
for view in views {
catalog_table_types
.entry(view.clone())
.or_insert(TableType::View);
names.push(view);
}
*self
.catalog_table_types
.write()
.unwrap_or_else(|e| e.into_inner()) = catalog_table_types;

if let Some(temp) = &self.temp_provider {
names.extend(temp.table_names());
}
Expand Down Expand Up @@ -540,6 +559,27 @@ impl SchemaProvider for PaimonSchemaProvider {
.await
}

async fn table_type(&self, name: &str) -> DFResult<Option<TableType>> {
if let Some(temp) = &self.temp_provider {
if let Some(table_type) = temp.table_type(name).await? {
return Ok(Some(table_type));
}
}

if let Some(table_type) = self
.catalog_table_types
.read()
.unwrap_or_else(|e| e.into_inner())
.get(name)
{
return Ok(Some(*table_type));
}

self.table(name)
.await
.map(|table| table.map(|table| table.table_type()))
}

fn table_exist(&self, name: &str) -> bool {
if let Some(temp) = &self.temp_provider {
if temp.table_exist(name) {
Expand Down
143 changes: 142 additions & 1 deletion crates/integrations/datafusion/tests/sql_context_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

//! SQL context integration tests for paimon-datafusion.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

use async_trait::async_trait;
Expand Down Expand Up @@ -50,6 +50,105 @@ async fn create_sql_context(catalog: Arc<FileSystemCatalog>) -> SQLContext {
ctx
}

struct MetadataListingCatalog {
get_table_calls: AtomicUsize,
}

impl MetadataListingCatalog {
fn new() -> Self {
Self {
get_table_calls: AtomicUsize::new(0),
}
}

fn get_table_calls(&self) -> usize {
self.get_table_calls.load(Ordering::SeqCst)
}
}

#[async_trait]
impl Catalog for MetadataListingCatalog {
async fn list_databases(&self) -> paimon::Result<Vec<String>> {
Ok(vec!["default".to_string()])
}

async fn create_database(
&self,
_name: &str,
_ignore_if_exists: bool,
_properties: std::collections::HashMap<String, String>,
) -> paimon::Result<()> {
Ok(())
}

async fn get_database(&self, name: &str) -> paimon::Result<paimon::catalog::Database> {
Ok(paimon::catalog::Database::new(
name.to_string(),
std::collections::HashMap::new(),
None,
))
}

async fn drop_database(
&self,
_name: &str,
_ignore_if_not_exists: bool,
_cascade: bool,
) -> paimon::Result<()> {
Ok(())
}

async fn get_table(&self, _identifier: &Identifier) -> paimon::Result<paimon::table::Table> {
self.get_table_calls.fetch_add(1, Ordering::SeqCst);
Err(paimon::Error::Unsupported {
message: "table loading is unavailable".to_string(),
})
}

async fn list_tables(&self, _database_name: &str) -> paimon::Result<Vec<String>> {
Ok(vec!["metadata_only".to_string()])
}

async fn list_views(&self, _database_name: &str) -> paimon::Result<Vec<String>> {
Ok(vec!["metadata_view".to_string()])
}

async fn create_table(
&self,
_identifier: &Identifier,
_creation: paimon::spec::Schema,
_ignore_if_exists: bool,
) -> paimon::Result<()> {
Ok(())
}

async fn drop_table(
&self,
_identifier: &Identifier,
_ignore_if_not_exists: bool,
) -> paimon::Result<()> {
Ok(())
}

async fn rename_table(
&self,
_from: &Identifier,
_to: &Identifier,
_ignore_if_not_exists: bool,
) -> paimon::Result<()> {
Ok(())
}

async fn alter_table(
&self,
_identifier: &Identifier,
_changes: Vec<SchemaChange>,
_ignore_if_not_exists: bool,
) -> paimon::Result<()> {
Ok(())
}
}

struct PartitionCatalog {
inner: Arc<FileSystemCatalog>,
fail_list_partitions: AtomicBool,
Expand Down Expand Up @@ -265,6 +364,48 @@ async fn test_show_tables_is_enabled() {
.expect("SHOW TABLES should execute");
}

#[tokio::test]
async fn test_show_tables_does_not_load_table_providers() {
let catalog = Arc::new(MetadataListingCatalog::new());
let mut sql_context = SQLContext::new();
sql_context
.register_catalog("paimon", catalog.clone())
.await
.unwrap();

let table_names = collect_string_column(&sql_context, "SHOW TABLES", "table_name").await;

assert!(table_names.contains(&"metadata_only".to_string()));
assert_eq!(catalog.get_table_calls(), 0);

assert!(sql_context
.sql("SELECT * FROM metadata_only")
.await
.is_err());
assert!(catalog.get_table_calls() > 0);
}

#[tokio::test]
async fn test_show_tables_preserves_catalog_view_type() {
let catalog = Arc::new(MetadataListingCatalog::new());
let mut sql_context = SQLContext::new();
sql_context
.register_catalog("paimon", catalog.clone())
.await
.unwrap();

let table_types = collect_string_column(
&sql_context,
"SELECT table_type FROM information_schema.tables \
WHERE table_schema = 'default' AND table_name = 'metadata_view'",
"table_type",
)
.await;

assert_eq!(table_types, vec!["VIEW"]);
assert_eq!(catalog.get_table_calls(), 0);
}

#[tokio::test]
async fn test_select_branch_table_reads_branch_snapshot() {
let (_tmp, catalog) = create_test_env();
Expand Down
Loading