micromegas_analytics/lakehouse/
retire_partition_by_file_udf.rs

1use anyhow::{Context, Result};
2use async_trait::async_trait;
3use datafusion::{
4    arrow::{
5        array::{Array, StringArray, StringBuilder},
6        datatypes::DataType,
7    },
8    common::internal_err,
9    error::DataFusionError,
10    logical_expr::{
11        ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
12        async_udf::AsyncScalarUDFImpl,
13    },
14};
15use micromegas_ingestion::data_lake_connection::DataLakeConnection;
16use micromegas_tracing::prelude::*;
17use sqlx::Row;
18use std::sync::Arc;
19
20use super::write_partition::add_file_for_cleanup;
21
22/// A scalar UDF that retires a single partition by its file path.
23///
24/// This function retires only the exact specified partition from the lakehouse.
25#[derive(Debug)]
26pub struct RetirePartitionByFile {
27    signature: Signature,
28    lake: Arc<DataLakeConnection>,
29}
30
31impl PartialEq for RetirePartitionByFile {
32    fn eq(&self, other: &Self) -> bool {
33        self.signature == other.signature
34    }
35}
36
37impl Eq for RetirePartitionByFile {}
38
39impl std::hash::Hash for RetirePartitionByFile {
40    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
41        self.signature.hash(state);
42    }
43}
44
45impl RetirePartitionByFile {
46    pub fn new(lake: Arc<DataLakeConnection>) -> Self {
47        Self {
48            signature: Signature::exact(vec![DataType::Utf8], Volatility::Volatile),
49            lake,
50        }
51    }
52
53    /// Retires a single partition by its file path within an existing transaction.
54    ///
55    /// # Arguments
56    /// * `transaction` - Database transaction to use
57    /// * `file_path` - The exact file path of the partition to retire
58    ///
59    /// # Returns
60    /// * `Ok(())` on successful retirement
61    /// * `Err(anyhow::Error)` with descriptive message for any failure
62    async fn retire_partition_in_transaction(
63        &self,
64        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
65        file_path: &str,
66    ) -> Result<()> {
67        // First, check if the partition exists and get its details
68        let partition_query = instrument_named!(
69            sqlx::query(
70                "SELECT file_path, file_size FROM lakehouse_partitions WHERE file_path = $1",
71            )
72            .bind(file_path)
73            .fetch_optional(&mut **transaction),
74            "sql_select_partition_by_file"
75        )
76        .await
77        .with_context(|| format!("querying partition {file_path}"))?;
78
79        let Some(partition_row) = partition_query else {
80            anyhow::bail!("Partition not found: {file_path}");
81        };
82
83        let file_size: i64 = partition_row.try_get("file_size")?;
84
85        // Add to temporary files for cleanup (expires in 1 hour)
86        add_file_for_cleanup(transaction, file_path, file_size).await?;
87
88        // Remove from active partitions
89        let delete_result = instrument_named!(
90            sqlx::query("DELETE FROM lakehouse_partitions WHERE file_path = $1")
91                .bind(file_path)
92                .execute(&mut **transaction),
93            "sql_delete_partition_by_file"
94        )
95        .await
96        .with_context(|| format!("deleting partition {file_path}"))?;
97
98        if delete_result.rows_affected() == 0 {
99            // This shouldn't happen since we checked existence above, but handle it gracefully
100            anyhow::bail!("Partition not found during deletion: {file_path}");
101        }
102
103        info!("Successfully retired partition: {}", file_path);
104        Ok(())
105    }
106}
107
108impl ScalarUDFImpl for RetirePartitionByFile {
109    fn name(&self) -> &str {
110        "retire_partition_by_file"
111    }
112
113    fn signature(&self) -> &Signature {
114        &self.signature
115    }
116
117    fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
118        Ok(DataType::Utf8)
119    }
120
121    fn invoke_with_args(
122        &self,
123        _args: ScalarFunctionArgs,
124    ) -> datafusion::error::Result<ColumnarValue> {
125        Err(DataFusionError::NotImplemented(
126            "retire_partition_by_file can only be called from async contexts".into(),
127        ))
128    }
129}
130
131#[async_trait]
132impl AsyncScalarUDFImpl for RetirePartitionByFile {
133    async fn invoke_async_with_args(
134        &self,
135        args: ScalarFunctionArgs,
136    ) -> datafusion::error::Result<ColumnarValue> {
137        let args = ColumnarValue::values_to_arrays(&args.args)?;
138        if args.len() != 1 {
139            return internal_err!("retire_partition_by_file expects exactly 1 argument: file_path");
140        }
141
142        let file_paths: &StringArray = args[0].as_any().downcast_ref::<_>().ok_or_else(|| {
143            DataFusionError::Execution("error casting file_path argument as StringArray".into())
144        })?;
145
146        let mut builder = StringBuilder::with_capacity(file_paths.len(), 64);
147
148        // Use a single transaction for the entire batch
149        let mut transaction =
150            self.lake.db_pool.begin().await.map_err(|e| {
151                DataFusionError::Execution(format!("Failed to begin transaction: {e}"))
152            })?;
153
154        let mut success_count = 0;
155        let mut has_errors = false;
156
157        // Process each file path in the batch within the same transaction
158        for index in 0..file_paths.len() {
159            if file_paths.is_null(index) {
160                builder.append_value("ERROR: file_path cannot be null");
161                has_errors = true;
162                continue;
163            }
164
165            let file_path = file_paths.value(index);
166
167            match self
168                .retire_partition_in_transaction(&mut transaction, file_path)
169                .await
170            {
171                Ok(()) => {
172                    success_count += 1;
173                    builder.append_value(format!("SUCCESS: Retired partition {file_path}"));
174                }
175                Err(e) => {
176                    error!("Failed to retire partition {}: {:?}", file_path, e);
177                    builder.append_value(format!("ERROR: {e:?}"));
178                    has_errors = true;
179                }
180            }
181        }
182
183        // Commit the transaction only if there were no errors
184        if has_errors {
185            if let Err(e) = transaction.rollback().await {
186                error!("Failed to rollback transaction after errors: {:?}", e);
187            }
188            info!("Rolled back transaction due to errors in batch retirement");
189        } else {
190            transaction.commit().await.map_err(|e| {
191                DataFusionError::Execution(format!("Failed to commit transaction: {e}"))
192            })?;
193            info!("Successfully retired {} partitions in batch", success_count);
194        }
195
196        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
197    }
198}
199
200/// Creates a user-defined function to retire a single partition by its file path.
201///
202/// This function ensures only the exact specified partition is removed from the lakehouse.
203///
204/// # Usage
205/// ```sql
206/// SELECT retire_partition_by_file('/path/to/partition.parquet') as result;
207/// ```
208///
209/// # Returns
210/// A string message indicating success or failure:
211/// - "SUCCESS: Retired partition <file_path>" on successful retirement
212/// - "ERROR: Partition not found: <file_path>" if the partition doesn't exist  
213/// - "ERROR: Database error: \<details\>" for any database-related failures
214pub fn make_retire_partition_by_file_udf(
215    lake: Arc<DataLakeConnection>,
216) -> datafusion::logical_expr::async_udf::AsyncScalarUDF {
217    datafusion::logical_expr::async_udf::AsyncScalarUDF::new(Arc::new(RetirePartitionByFile::new(
218        lake,
219    )))
220}
micromegas_analytics/lakehouse/retire_partition_by_file_udf.rs

micromegas_analytics/lakehouse/
retire_partition_by_file_udf.rs
