diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 83737892..8ce60dbb 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -61,21 +61,10 @@ class Table: def identifier(self) -> str: ... def location(self) -> str: ... def schema(self) -> TableSchema: ... - def new_read_builder(self, options: Optional[Dict[str, str]] = None) -> ReadBuilder: ... - def new_write_builder(self) -> "WriteBuilder": ... - -class CommitMessage: ... - -class TableWrite: - def write_arrow(self, batch: pyarrow.RecordBatch) -> None: ... - def prepare_commit(self) -> List[CommitMessage]: ... - -class TableCommit: - def commit(self, messages: Sequence[CommitMessage]) -> None: ... - -class WriteBuilder: - def new_write(self) -> TableWrite: ... - def new_commit(self) -> TableCommit: ... + def new_read_builder(self) -> ReadBuilder: ... + def expire_snapshots(self, older_than_ms: int) -> int: ... + def remove_orphan_files(self) -> int: ... + def drop_partition(self, partition: Dict[str, Any]) -> None: ... class PaimonCatalog: def __init__(self, catalog_options: Dict[str, str]) -> None: ... diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs index a59ca8fe..eb47912a 100644 --- a/bindings/python/src/table.rs +++ b/bindings/python/src/table.rs @@ -14,12 +14,17 @@ // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. - +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use std::collections::HashMap; use std::sync::Arc; -use pyo3::prelude::*; -use pyo3::types::PyDict; +use paimon::spec::Datum; +use paimon::table::SnapshotManager; +use paimon_datafusion::runtime::runtime; +use crate::error::to_py_err; +use crate::predicate::py_to_datum; use crate::read::PyReadBuilder; use crate::schema::PyTableSchema; use crate::write::PyWriteBuilder; @@ -68,4 +73,62 @@ impl PyTable { fn new_write_builder(&self) -> PyWriteBuilder { PyWriteBuilder::new(Arc::clone(&self.inner)) } + fn expire_snapshots(&self, py: Python<'_>, older_than_ms: i64) -> PyResult { + let rt = runtime(); + py.detach(|| { + rt.block_on(async { + let snapshot_manager = SnapshotManager::new( + self.inner.file_io().clone(), + self.inner.location().to_string(), + ); + snapshot_manager + .expire_snapshots_earlier_than(older_than_ms) + .await + .map_err(to_py_err) + }) + }) + } + + fn remove_orphan_files(&self, py: Python<'_>) -> PyResult { + let rt = runtime(); + py.detach(|| { + rt.block_on(async { + let snapshot_manager = SnapshotManager::new( + self.inner.file_io().clone(), + self.inner.location().to_string(), + ); + snapshot_manager + .remove_orphan_files() + .await + .map_err(to_py_err) + }) + }) + } + + fn drop_partition(&self, partition: HashMap>) -> PyResult<()> { + let partition_fields = self.inner.schema().partition_fields(); + let mut spec: HashMap> = HashMap::with_capacity(partition.len()); + for (k, v) in &partition { + let datum = if v.is_none() { + None + } else { + let field = partition_fields + .iter() + .find(|f| f.name() == k) + .ok_or_else(|| { + PyValueError::new_err(format!("Partition field {} not found in schema", k)) + })?; + Some(py_to_datum(v, field.data_type())?) + }; + spec.insert(k.clone(), datum); + } + runtime().block_on(async { + let commit = self.inner.new_write_builder().new_commit(); + commit.drop_partitions(vec![spec]).await.map_err(to_py_err) + }) + } + + fn trigger_compaction(&self, _full_compact: bool) -> PyResult<()> { + todo!() + } } diff --git a/crates/paimon/src/table/snapshot_manager.rs b/crates/paimon/src/table/snapshot_manager.rs index 73effe69..1b5706c1 100644 --- a/crates/paimon/src/table/snapshot_manager.rs +++ b/crates/paimon/src/table/snapshot_manager.rs @@ -345,6 +345,60 @@ impl SnapshotManager { Ok(result) } + pub async fn expire_snapshots(&self, snapshot_ids: &[i64]) -> crate::Result { + let mut deleted_count = 0; + for &snapshot_id in snapshot_ids { + self.delete_snapshot(snapshot_id).await?; + deleted_count += 1; + } + Ok(deleted_count) + } + + // Expires snapshots whose commit time is earlier than target timestamp_millis. + // Returns the number of snapshots deleted. + pub async fn expire_snapshots_earlier_than(&self, timestamp_millis: i64) -> crate::Result { + let snapshot_ids = self.list_all_ids().await?; + let mut to_delete = Vec::new(); + for snapshot_id in snapshot_ids { + let snapshot = self.get_snapshot(snapshot_id).await?; + if (snapshot.time_millis() as i64) < timestamp_millis { + to_delete.push(snapshot_id); + } else { + break; + } + } + self.expire_snapshots(&to_delete).await?; + Ok(to_delete.len() as i64) + } + + // Remove files which are not used by any snapshots and are not referenced by + // manifest lists. + pub async fn remove_orphan_files(&self) -> crate::Result { + let snapshot_ids = self.list_all_ids().await?; + let mut manifest_files = std::collections::HashSet::new(); + for snap_id in snapshot_ids { + let snapshot = self.get_snapshot(snap_id).await?; + manifest_files.insert(snapshot.base_manifest_list().to_string()); + manifest_files.insert(snapshot.delta_manifest_list().to_string()); + } + let manifest_dir = self.manifest_dir(); + let statuses = self.file_io.list_status(&manifest_dir).await?; + let mut deleted_count = 0; + for status in statuses { + if status.is_dir { + continue; + } + let name = status.path.rsplit('/').next().unwrap_or(&status.path); + let manifest_path = format!("{}/{}", manifest_dir, name); + if !manifest_files.contains(&manifest_path) { + self.file_io.delete_file(&manifest_path).await?; + deleted_count += 1; + } + } + + Ok(deleted_count) + } + /// Returns the snapshot whose commit time is earlier than or equal to the given /// `timestamp_millis`. If no such snapshot exists, returns None. ///