Skip to content

Commit 2b58d25

Browse files
JingsongLiJingsongLi
authored
refactor(datafusion): convert files_size from table functions to system tables (#325)
These are table metadata and fit better as system tables (table$physical_files_size) rather than UDTFs (physical_files_size('table')), consistent with $partitions, $snapshots, etc.
1 parent ab68b26 commit 2b58d25

5 files changed

Lines changed: 78 additions & 238 deletions

File tree

‎crates/integrations/datafusion/src/lib.rs‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,10 +43,8 @@ mod filter_pushdown;
4343
#[cfg(feature = "fulltext")]
4444
mod full_text_search;
4545
mod merge_into;
46-
mod physical_files_size;
4746
mod physical_plan;
4847
mod procedures;
49-
mod referenced_files_size;
5048
mod relation_planner;
5149
pub mod runtime;
5250
mod sql_context;
@@ -69,9 +67,7 @@ pub use catalog::{PaimonCatalogProvider, PaimonSchemaProvider};
6967
pub use error::to_datafusion_error;
7068
#[cfg(feature = "fulltext")]
7169
pub use full_text_search::{register_full_text_search, FullTextSearchFunction};
72-
pub use physical_files_size::{register_physical_files_size, PhysicalFilesSizeFunction};
7370
pub use physical_plan::PaimonTableScan;
74-
pub use referenced_files_size::{register_referenced_files_size, ReferencedFilesSizeFunction};
7571
pub use relation_planner::PaimonRelationPlanner;
7672
pub use sql_context::SQLContext;
7773
pub use table::PaimonTableProvider;

‎crates/integrations/datafusion/src/system_tables/mod.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ mod branches;
3333
mod manifests;
3434
mod options;
3535
mod partitions;
36+
mod physical_files_size;
37+
mod referenced_files_size;
3638
mod row_string_cast;
3739
mod schemas;
3840
mod snapshots;
@@ -47,6 +49,8 @@ const TABLES: &[(&str, Builder)] = &[
4749
("branches", branches::build),
4850
("manifests", manifests::build),
4951
("options", options::build),
52+
("physical_files_size", physical_files_size::build),
53+
("referenced_files_size", referenced_files_size::build),
5054
("schemas", schemas::build),
5155
("snapshots", snapshots::build),
5256
("tags", tags::build),
@@ -57,6 +61,8 @@ const SYSTEM_TABLE_NAMES: &[&str] = &[
5761
"manifests",
5862
"options",
5963
"partitions",
64+
"physical_files_size",
65+
"referenced_files_size",
6066
"schemas",
6167
"snapshots",
6268
"tags",

crates/integrations/datafusion/src/physical_files_size.rs renamed to crates/integrations/datafusion/src/system_tables/physical_files_size.rs

Lines changed: 6 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -15,89 +15,27 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
//! Table function that computes the total physical file sizes in the table directory.
19-
//!
20-
//! Usage: `SELECT * FROM physical_files_size('db.table_name')`
18+
//! Mirrors Java [PhysicalFilesSizeTable](https://github.com/apache/paimon/blob/release-1.4/paimon-core/src/main/java/org/apache/paimon/table/system).
2119
2220
use std::any::Any;
23-
use std::fmt::Debug;
2421
use std::sync::{Arc, OnceLock};
2522

2623
use async_trait::async_trait;
2724
use datafusion::arrow::array::{Int64Array, RecordBatch};
2825
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
2926
use datafusion::catalog::Session;
30-
use datafusion::catalog::TableFunctionImpl;
3127
use datafusion::datasource::memory::MemorySourceConfig;
3228
use datafusion::datasource::{TableProvider, TableType};
3329
use datafusion::error::Result as DFResult;
3430
use datafusion::logical_expr::Expr;
3531
use datafusion::physical_plan::ExecutionPlan;
36-
use datafusion::prelude::SessionContext;
37-
use paimon::catalog::Catalog;
3832
use paimon::table::referenced_files::{collect_physical_files_summary, PhysicalFilesSummary};
3933
use paimon::table::Table;
4034

4135
use crate::error::to_datafusion_error;
42-
use crate::runtime::{await_with_runtime, block_on_with_runtime};
43-
use crate::table_function_args::{extract_string_literal, parse_table_identifier};
4436

45-
const FUNCTION_NAME: &str = "physical_files_size";
46-
47-
pub fn register_physical_files_size(
48-
ctx: &SessionContext,
49-
catalog: Arc<dyn Catalog>,
50-
default_database: &str,
51-
) {
52-
ctx.register_udtf(
53-
FUNCTION_NAME,
54-
Arc::new(PhysicalFilesSizeFunction::new(catalog, default_database)),
55-
);
56-
}
57-
58-
pub struct PhysicalFilesSizeFunction {
59-
catalog: Arc<dyn Catalog>,
60-
default_database: String,
61-
}
62-
63-
impl Debug for PhysicalFilesSizeFunction {
64-
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65-
f.debug_struct("PhysicalFilesSizeFunction")
66-
.field("default_database", &self.default_database)
67-
.finish()
68-
}
69-
}
70-
71-
impl PhysicalFilesSizeFunction {
72-
pub fn new(catalog: Arc<dyn Catalog>, default_database: &str) -> Self {
73-
Self {
74-
catalog,
75-
default_database: default_database.to_string(),
76-
}
77-
}
78-
}
79-
80-
impl TableFunctionImpl for PhysicalFilesSizeFunction {
81-
fn call(&self, args: &[Expr]) -> DFResult<Arc<dyn TableProvider>> {
82-
if args.len() != 1 {
83-
return Err(datafusion::error::DataFusionError::Plan(
84-
"physical_files_size requires 1 argument: (table_name)".to_string(),
85-
));
86-
}
87-
88-
let table_name = extract_string_literal(FUNCTION_NAME, &args[0], "table_name")?;
89-
let identifier =
90-
parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?;
91-
92-
let catalog = Arc::clone(&self.catalog);
93-
let table = block_on_with_runtime(
94-
async move { catalog.get_table(&identifier).await },
95-
"physical_files_size: catalog access thread panicked",
96-
)
97-
.map_err(to_datafusion_error)?;
98-
99-
Ok(Arc::new(PhysicalFilesSizeTableProvider { table }))
100-
}
37+
pub(super) fn build(table: Table) -> DFResult<Arc<dyn TableProvider>> {
38+
Ok(Arc::new(PhysicalFilesSizeTable { table }))
10139
}
10240

10341
fn output_schema() -> SchemaRef {
@@ -117,12 +55,12 @@ fn output_schema() -> SchemaRef {
11755
}
11856

11957
#[derive(Debug)]
120-
struct PhysicalFilesSizeTableProvider {
58+
struct PhysicalFilesSizeTable {
12159
table: Table,
12260
}
12361

12462
#[async_trait]
125-
impl TableProvider for PhysicalFilesSizeTableProvider {
63+
impl TableProvider for PhysicalFilesSizeTable {
12664
fn as_any(&self) -> &dyn Any {
12765
self
12866
}
@@ -143,7 +81,7 @@ impl TableProvider for PhysicalFilesSizeTableProvider {
14381
_limit: Option<usize>,
14482
) -> DFResult<Arc<dyn ExecutionPlan>> {
14583
let table = self.table.clone();
146-
let summary = await_with_runtime(async move {
84+
let summary = crate::runtime::await_with_runtime(async move {
14785
collect_physical_files_summary(table.file_io(), table.location()).await
14886
})
14987
.await

crates/integrations/datafusion/src/referenced_files_size.rs renamed to crates/integrations/datafusion/src/system_tables/referenced_files_size.rs

Lines changed: 6 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -15,89 +15,27 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
//! Table function that computes per-snapshot referenced file size summaries.
19-
//!
20-
//! Usage: `SELECT * FROM referenced_files_size('db.table_name')`
18+
//! Mirrors Java [ReferencedFilesSizeTable](https://github.com/apache/paimon/blob/release-1.4/paimon-core/src/main/java/org/apache/paimon/table/system).
2119
2220
use std::any::Any;
23-
use std::fmt::Debug;
2421
use std::sync::{Arc, OnceLock};
2522

2623
use async_trait::async_trait;
2724
use datafusion::arrow::array::{Int64Array, RecordBatch, StringArray};
2825
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
2926
use datafusion::catalog::Session;
30-
use datafusion::catalog::TableFunctionImpl;
3127
use datafusion::datasource::memory::MemorySourceConfig;
3228
use datafusion::datasource::{TableProvider, TableType};
3329
use datafusion::error::Result as DFResult;
3430
use datafusion::logical_expr::Expr;
3531
use datafusion::physical_plan::ExecutionPlan;
36-
use datafusion::prelude::SessionContext;
37-
use paimon::catalog::Catalog;
3832
use paimon::table::referenced_files::{collect_referenced_files_summary, ReferencedFilesSummary};
3933
use paimon::table::Table;
4034

4135
use crate::error::to_datafusion_error;
42-
use crate::runtime::{await_with_runtime, block_on_with_runtime};
43-
use crate::table_function_args::{extract_string_literal, parse_table_identifier};
44-
45-
const FUNCTION_NAME: &str = "referenced_files_size";
46-
47-
pub fn register_referenced_files_size(
48-
ctx: &SessionContext,
49-
catalog: Arc<dyn Catalog>,
50-
default_database: &str,
51-
) {
52-
ctx.register_udtf(
53-
FUNCTION_NAME,
54-
Arc::new(ReferencedFilesSizeFunction::new(catalog, default_database)),
55-
);
56-
}
57-
58-
pub struct ReferencedFilesSizeFunction {
59-
catalog: Arc<dyn Catalog>,
60-
default_database: String,
61-
}
62-
63-
impl Debug for ReferencedFilesSizeFunction {
64-
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65-
f.debug_struct("ReferencedFilesSizeFunction")
66-
.field("default_database", &self.default_database)
67-
.finish()
68-
}
69-
}
7036

71-
impl ReferencedFilesSizeFunction {
72-
pub fn new(catalog: Arc<dyn Catalog>, default_database: &str) -> Self {
73-
Self {
74-
catalog,
75-
default_database: default_database.to_string(),
76-
}
77-
}
78-
}
79-
80-
impl TableFunctionImpl for ReferencedFilesSizeFunction {
81-
fn call(&self, args: &[Expr]) -> DFResult<Arc<dyn TableProvider>> {
82-
if args.len() != 1 {
83-
return Err(datafusion::error::DataFusionError::Plan(
84-
"referenced_files_size requires 1 argument: (table_name)".to_string(),
85-
));
86-
}
87-
88-
let table_name = extract_string_literal(FUNCTION_NAME, &args[0], "table_name")?;
89-
let identifier =
90-
parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?;
91-
92-
let catalog = Arc::clone(&self.catalog);
93-
let table = block_on_with_runtime(
94-
async move { catalog.get_table(&identifier).await },
95-
"referenced_files_size: catalog access thread panicked",
96-
)
97-
.map_err(to_datafusion_error)?;
98-
99-
Ok(Arc::new(ReferencedFilesSizeTableProvider { table }))
100-
}
37+
pub(super) fn build(table: Table) -> DFResult<Arc<dyn TableProvider>> {
38+
Ok(Arc::new(ReferencedFilesSizeTable { table }))
10139
}
10240

10341
fn output_schema() -> SchemaRef {
@@ -118,12 +56,12 @@ fn output_schema() -> SchemaRef {
11856
}
11957

12058
#[derive(Debug)]
121-
struct ReferencedFilesSizeTableProvider {
59+
struct ReferencedFilesSizeTable {
12260
table: Table,
12361
}
12462

12563
#[async_trait]
126-
impl TableProvider for ReferencedFilesSizeTableProvider {
64+
impl TableProvider for ReferencedFilesSizeTable {
12765
fn as_any(&self) -> &dyn Any {
12866
self
12967
}
@@ -144,7 +82,7 @@ impl TableProvider for ReferencedFilesSizeTableProvider {
14482
_limit: Option<usize>,
14583
) -> DFResult<Arc<dyn ExecutionPlan>> {
14684
let table = self.table.clone();
147-
let summaries = await_with_runtime(async move {
85+
let summaries = crate::runtime::await_with_runtime(async move {
14886
let schema = table.schema();
14987
let partition_keys = schema.partition_keys();
15088
let partition_fields = schema.partition_fields();

0 commit comments

Comments
 (0)