use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, OnceLock};
use async_trait::async_trait;
use datafusion::arrow::array::{
BooleanArray, Int32Array, Int64Array, RecordBatch, StringArray, TimestampMillisecondArray,
};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
use datafusion::catalog::Session;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::{TableProvider, TableType};
use datafusion::error::{DataFusionError, Result as DFResult};
use datafusion::logical_expr::Expr;
use datafusion::physical_plan::ExecutionPlan;
use paimon::catalog::{list_partitions_from_file_system, Catalog, Identifier};
use paimon::spec::{CoreOptions, Partition};
use paimon::table::Table;
use crate::error::to_datafusion_error;
pub(super) fn build(
catalog: Arc<dyn Catalog>,
identifier: Identifier,
table: Table,
) -> DFResult<Arc<dyn TableProvider>> {
let table_schema = table.schema();
let partition_keys = table_schema.partition_keys().to_vec();
let default_partition_name = CoreOptions::new(table_schema.options())
.partition_default_name()
.to_string();
Ok(Arc::new(PartitionsTable {
catalog,
identifier,
table,
partition_keys,
default_partition_name,
}))
}
fn partitions_schema() -> SchemaRef {
static SCHEMA: OnceLock<SchemaRef> = OnceLock::new();
SCHEMA
.get_or_init(|| {
Arc::new(Schema::new(vec![
Field::new("partition", DataType::Utf8, true),
Field::new("record_count", DataType::Int64, false),
Field::new("file_size_in_bytes", DataType::Int64, false),
Field::new("file_count", DataType::Int64, false),
Field::new(
"last_update_time",
DataType::Timestamp(TimeUnit::Millisecond, None),
true,
),
Field::new(
"created_at",
DataType::Timestamp(TimeUnit::Millisecond, None),
true,
),
Field::new("created_by", DataType::Utf8, true),
Field::new("updated_by", DataType::Utf8, true),
Field::new("options", DataType::Utf8, true),
Field::new("total_buckets", DataType::Int32, false),
Field::new("done", DataType::Boolean, false),
]))
})
.clone()
}
struct PartitionsTable {
catalog: Arc<dyn Catalog>,
identifier: Identifier,
table: Table,
partition_keys: Vec<String>,
default_partition_name: String,
}
impl std::fmt::Debug for PartitionsTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PartitionsTable")
.field("identifier", &self.identifier)
.finish()
}
}
#[async_trait]
impl TableProvider for PartitionsTable {
fn schema(&self) -> SchemaRef {
partitions_schema()
}
fn table_type(&self) -> TableType {
TableType::View
}
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
_filters: &[Expr],
_limit: Option<usize>,
) -> DFResult<Arc<dyn ExecutionPlan>> {
let table = self.table.clone();
let partitions = if table.travel_snapshot().is_some() {
crate::runtime::await_with_runtime(async move {
list_partitions_from_file_system(&table).await
})
.await
} else {
let catalog = self.catalog.clone();
let identifier = if table.is_main_branch() {
self.identifier.clone()
} else {
Identifier::new(
self.identifier.database(),
format!("{}$branch_{}", self.identifier.object(), table.branch()),
)
};
crate::runtime::await_with_runtime(async move {
match catalog.list_partitions(&identifier).await {
Ok(partitions) => Ok(partitions),
Err(_) => list_partitions_from_file_system(&table).await,
}
})
.await
}
.map_err(to_datafusion_error)?;
let mut rows: Vec<(String, Partition)> = partitions
.into_iter()
.map(|p| {
let s = format_partition_string(
&p.spec,
&self.partition_keys,
&self.default_partition_name,
);
(s, p)
})
.collect();
rows.sort_by(|a, b| a.0.cmp(&b.0));
let n = rows.len();
let mut partition_strings: Vec<Option<String>> = Vec::with_capacity(n);
let mut record_counts = Vec::with_capacity(n);
let mut file_sizes = Vec::with_capacity(n);
let mut file_counts = Vec::with_capacity(n);
let mut last_update_times: Vec<Option<i64>> = Vec::with_capacity(n);
let mut created_ats: Vec<Option<i64>> = Vec::with_capacity(n);
let mut created_bys: Vec<Option<String>> = Vec::with_capacity(n);
let mut updated_bys: Vec<Option<String>> = Vec::with_capacity(n);
let mut options_jsons: Vec<Option<String>> = Vec::with_capacity(n);
let mut total_buckets = Vec::with_capacity(n);
let mut dones = Vec::with_capacity(n);
for (s, p) in rows {
partition_strings.push(Some(s));
record_counts.push(p.record_count);
file_sizes.push(p.file_size_in_bytes);
file_counts.push(p.file_count);
last_update_times.push(if p.last_file_creation_time > 0 {
Some(p.last_file_creation_time)
} else {
None
});
created_ats.push(p.created_at);
created_bys.push(p.created_by);
updated_bys.push(p.updated_by);
options_jsons.push(
p.options
.as_ref()
.map(|m| {
let sorted: BTreeMap<&String, &String> = m.iter().collect();
serde_json::to_string(&sorted)
.map_err(|e| DataFusionError::External(Box::new(e)))
})
.transpose()?,
);
total_buckets.push(p.total_buckets);
dones.push(p.done);
}
let schema = partitions_schema();
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(StringArray::from(partition_strings)),
Arc::new(Int64Array::from(record_counts)),
Arc::new(Int64Array::from(file_sizes)),
Arc::new(Int64Array::from(file_counts)),
Arc::new(TimestampMillisecondArray::from(last_update_times)),
Arc::new(TimestampMillisecondArray::from(created_ats)),
Arc::new(StringArray::from(created_bys)),
Arc::new(StringArray::from(updated_bys)),
Arc::new(StringArray::from(options_jsons)),
Arc::new(Int32Array::from(total_buckets)),
Arc::new(BooleanArray::from(dones)),
],
)?;
Ok(MemorySourceConfig::try_new_exec(
&[vec![batch]],
schema,
projection.cloned(),
)?)
}
}
fn format_partition_string(
spec: &HashMap<String, String>,
partition_keys: &[String],
default_partition_name: &str,
) -> String {
partition_keys
.iter()
.map(|k| {
let v = spec
.get(k)
.map(String::as_str)
.unwrap_or(default_partition_name);
format!("{k}={v}")
})
.collect::<Vec<_>>()
.join("/")
}