Add support for manually optimizing a table

crates/modelardb_embedded/bindings/c/modelardb_embedded.h

@@ -199,6 +199,12 @@ int modelardb_embedded_vacuum(void* maybe_operations_ptr,
                               const char* table_name_ptr,
                               const char* retention_period_in_seconds_ptr);
 
+// Optimize the table by compacting its many small files into fewer larger files.
+int modelardb_embedded_optimize(void* maybe_operations_ptr,
+                                bool is_data_folder,
+                                const char* table_name_ptr,
+                                const char* target_size_in_bytes_ptr);
+
 // Return a human-readable representation of the last error on this thread.
 const char* modelardb_embedded_error();
 

crates/modelardb_embedded/bindings/python/modelardb/operations.py

@@ -682,6 +682,37 @@ def vacuum(self, table_name: str, retention_period_in_seconds: None | int = None
         )
         self.__check_return_code_and_raise_error(return_code)
 
+    def optimize(self, table_name: str, target_size_in_bytes: None | int = None):
+        """Optimize the table with `table_name`.
+
+        :param table_name: The name of the table to optimize.
+        :type table_name: str
+        :param target_size_in_bytes: The target file size in bytes. Many small files are compacted
+         into fewer larger files of approximately this size. If `None`, the default target size of
+         64 MiB is used.
+        :type target_size_in_bytes: int, optional
+        :raises ValueError: If incorrect arguments are provided.
+        """
+        table_name_ptr = ffi.new("char[]", bytes(table_name, "UTF-8"))
+
+        if target_size_in_bytes is not None:
+            # Convert the target size to a string to avoid issues with converting an int to a C type that uses
+            # an inconsistent amount of bits across platforms and then converting that to a 64-bit integer in Rust.
+            # The string is converted directly to an unsigned 64-bit integer in Rust.
+            target_size_in_bytes_ptr = ffi.new(
+                "char[]", bytes(str(target_size_in_bytes), "UTF-8")
+            )
+        else:
+            target_size_in_bytes_ptr = ffi.NULL
+
+        return_code = self.__library.modelardb_embedded_optimize(
+            self.__operations_ptr,
+            self.__is_data_folder,
+            table_name_ptr,
+            target_size_in_bytes_ptr,
+        )
+        self.__check_return_code_and_raise_error(return_code)
+
     def __check_return_code_and_raise_error(self, return_code: int):
         """Raises an appropriate exception based on the return code.
 

crates/modelardb_embedded/bindings/python/tests/test_operations.py

@@ -12,6 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Tests for the data folder operations exposed through the Python bindings.
+
+Test methods suffixed with ``_error`` only verify that errors from the
+underlying Rust implementation are propagated through the bindings and surface
+as Python exceptions. The errors themselves are already tested in the Rust
+implementation, so the specific error raised does not matter.
+"""
+
 import datetime
 import os
 import unittest
@@ -486,6 +494,41 @@ def test_data_folder_vacuum_error(self):
             error_message = f"Invalid Argument Error: Table with name '{MISSING_TABLE_NAME}' does not exist."
             self.assertEqual(error_message, str(context.exception))
 
+    def test_data_folder_optimize(self):
+        with TemporaryDirectory() as temp_dir:
+            data_folder = Operations.open_local(temp_dir)
+            create_tables_in_data_folder(data_folder)
+
+            # Each write is a separate commit, so four writes produce multiple small files.
+            for _ in range(4):
+                data_folder.write(
+                    TIME_SERIES_TABLE_NAME,
+                    unsorted_time_series_table_data_without_generated_column(),
+                )
+
+            folder_path = os.path.join(temp_dir, "tables", TIME_SERIES_TABLE_NAME, "field_column=2")
+            file_count = len(os.listdir(folder_path))
+            self.assertEqual(file_count, 4)
+
+            data_folder.optimize(TIME_SERIES_TABLE_NAME)
+
+            # Vacuum to remove the compacted files.
+            data_folder.vacuum(TIME_SERIES_TABLE_NAME, retention_period_in_seconds=0)
+
+            # The small files should be compacted into a single active file.
+            file_count = len(os.listdir(folder_path))
+            self.assertEqual(file_count, 1)
+
+    def test_data_folder_optimize_error(self):
+        with TemporaryDirectory() as temp_dir:
+            data_folder = Operations.open_local(temp_dir)
+
+            with self.assertRaises(RuntimeError) as context:
+                data_folder.optimize(MISSING_TABLE_NAME)
+
+            error_message = f"Invalid Argument Error: Table with name '{MISSING_TABLE_NAME}' does not exist."
+            self.assertEqual(error_message, str(context.exception))
+
 
 def create_tables_in_data_folder(data_folder: Operations):
     table_type = NormalTable(normal_table_schema())

crates/modelardb_embedded/src/capi.rs

@@ -965,9 +965,9 @@ unsafe fn drop(
 /// `maybe_operations_ptr` by deleting stale files that are older than `retention_period_in_seconds_ptr`
 /// seconds. Assumes `maybe_operations_ptr` points to a [`DataFolder`] or [`Client`];
 /// `table_name_ptr` points to a valid C string; and `retention_period_in_seconds_ptr` points to a
-/// valid C string. A C string is used for the retention period to avoid issues with different
-/// platforms using an inconsistent amount of bits for integer types. The string is converted
-/// directly to an unsigned 64-bit integer in Rust.
+/// valid C string, or is null to use the default retention period. A C string is used for the
+/// retention period to avoid issues with different platforms using an inconsistent amount of bits
+/// for integer types. The string is converted directly to an unsigned 64-bit integer in Rust.
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn modelardb_embedded_vacuum(
     maybe_operations_ptr: *mut c_void,
@@ -1013,6 +1013,57 @@ unsafe fn vacuum(
     TOKIO_RUNTIME.block_on(modelardb.vacuum(table_name, maybe_retention_period_in_seconds))
 }
 
+/// Optimizes the table with the name in `table_name_ptr` in the [`DataFolder`] or [`Client`] in
+/// `maybe_operations_ptr` by compacting its many small files into fewer larger files of
+/// approximately `target_size_in_bytes_ptr` bytes. Assumes `maybe_operations_ptr` points to a
+/// [`DataFolder`] or [`Client`]; `table_name_ptr` points to a valid C string; and
+/// `target_size_in_bytes_ptr` points to a valid C string, or is null to use the default target
+/// size. A C string is used for the target size to avoid issues with different platforms using an
+/// inconsistent amount of bits for integer types. The string is converted directly to an unsigned
+/// 64-bit integer in Rust.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn modelardb_embedded_optimize(
+    maybe_operations_ptr: *mut c_void,
+    is_data_folder: bool,
+    table_name_ptr: *const c_char,
+    target_size_in_bytes_ptr: *const c_char,
+) -> c_int {
+    let maybe_unit = unsafe {
+        optimize(
+            maybe_operations_ptr,
+            is_data_folder,
+            table_name_ptr,
+            target_size_in_bytes_ptr,
+        )
+    };
+    set_error_and_return_code(maybe_unit)
+}
+
+/// See documentation for [`modelardb_embedded_optimize()`].
+unsafe fn optimize(
+    maybe_operations_ptr: *mut c_void,
+    is_data_folder: bool,
+    table_name_ptr: *const c_char,
+    target_size_in_bytes_ptr: *const c_char,
+) -> Result<()> {
+    let modelardb = unsafe { c_void_to_operations(maybe_operations_ptr, is_data_folder)? };
+    let table_name = unsafe { c_char_ptr_to_str(table_name_ptr)? };
+    let maybe_target_size_in_bytes_str =
+        unsafe { c_char_ptr_to_maybe_str(target_size_in_bytes_ptr)? };
+
+    let maybe_target_size_in_bytes = maybe_target_size_in_bytes_str
+        .map(|target_size_in_bytes_str| {
+            target_size_in_bytes_str.parse::<u64>().map_err(|error| {
+                ModelarDbEmbeddedError::InvalidArgument(format!(
+                    "Target size is not a valid u64: {error}"
+                ))
+            })
+        })
+        .transpose()?;
+
+    TOKIO_RUNTIME.block_on(modelardb.optimize(table_name, maybe_target_size_in_bytes))
+}
+
 /// Return a read-only [`*const c_char`] with a human-readable representation of the last error the
 /// current thread encountered. The lifetime of the returned [`*const c_char`] ends when
 /// [`modelardb_embedded_error()`] is called again. If no errors have occurred, a zero-initialized

crates/modelardb_embedded/src/operations/client.rs

@@ -310,4 +310,25 @@ impl Operations for Client {
 
         Ok(())
     }
+
+    /// Optimize the table with the name in `table_name` by compacting its many small files into
+    /// fewer larger files of approximately `maybe_target_size_in_bytes` bytes. If a target size is
+    /// not given, the default target size of 64 MiB is used. If the table does not exist, the table
+    /// could not be optimized, or the target size is zero, [`ModelarDbEmbeddedError`] is returned.
+    async fn optimize(
+        &mut self,
+        table_name: &str,
+        maybe_target_size_in_bytes: Option<u64>,
+    ) -> Result<()> {
+        let sql = if let Some(target_size_in_bytes) = maybe_target_size_in_bytes {
+            format!("OPTIMIZE {table_name} TARGET {target_size_in_bytes}")
+        } else {
+            format!("OPTIMIZE {table_name}")
+        };
+
+        let ticket = Ticket::new(sql);
+        self.flight_client.do_get(ticket).await?;
+
+        Ok(())
+    }
 }

crates/modelardb_embedded/src/operations/data_folder.rs

@@ -539,6 +539,26 @@ impl Operations for DataFolder {
             )))
         }
     }
+
+    /// Optimize the table with the name in `table_name` by compacting its many small files into
+    /// fewer larger files of approximately `maybe_target_size_in_bytes` bytes. If a target size is
+    /// not given, the default target size of 64 MiB is used. If the table does not exist, the table
+    /// could not be optimized, or the target size is zero, [`ModelarDbEmbeddedError`] is returned.
+    async fn optimize(
+        &mut self,
+        table_name: &str,
+        maybe_target_size_in_bytes: Option<u64>,
+    ) -> Result<()> {
+        if self.tables().await?.contains(&table_name.to_owned()) {
+            self.optimize_table(table_name, maybe_target_size_in_bytes)
+                .await
+                .map_err(|error| error.into())
+        } else {
+            Err(ModelarDbEmbeddedError::InvalidArgument(format!(
+                "Table with name '{table_name}' does not exist."
+            )))
+        }
+    }
 }
 
 /// Compare `source_schema` and `target_schema` and return [`true`] if they have the same number of
@@ -2276,6 +2296,75 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_optimize_normal_table() {
+        let (_temp_dir, mut data_folder) = create_data_folder_with_normal_table().await;
+
+        // Each write is a separate commit, so four writes produce four small files.
+        for _ in 0..4 {
+            data_folder
+                .write(NORMAL_TABLE_NAME, normal_table_data())
+                .await
+                .unwrap();
+        }
+
+        let delta_table = data_folder.delta_table(NORMAL_TABLE_NAME).await.unwrap();
+        assert_eq!(delta_table.get_file_uris().unwrap().count(), 4);
+
+        data_folder.optimize(NORMAL_TABLE_NAME, None).await.unwrap();
+
+        // The small files should be compacted into a single active file.
+        let delta_table = data_folder.delta_table(NORMAL_TABLE_NAME).await.unwrap();
+        assert_eq!(delta_table.get_file_uris().unwrap().count(), 1);
+    }
+
+    #[tokio::test]
+    async fn test_optimize_time_series_table() {
+        let (_temp_dir, mut data_folder) = create_data_folder_with_time_series_table().await;
+
+        // Each write commits one file per field column partition, so four writes to the two
+        // partitions produce eight small files.
+        for _ in 0..4 {
+            data_folder
+                .write(TIME_SERIES_TABLE_NAME, time_series_table_data())
+                .await
+                .unwrap();
+        }
+
+        let delta_table = data_folder
+            .delta_table(TIME_SERIES_TABLE_NAME)
+            .await
+            .unwrap();
+        assert_eq!(delta_table.get_file_uris().unwrap().count(), 8);
+
+        data_folder
+            .optimize(TIME_SERIES_TABLE_NAME, None)
+            .await
+            .unwrap();
+
+        // The files in each of the two partitions should be compacted into a single active file.
+        let delta_table = data_folder
+            .delta_table(TIME_SERIES_TABLE_NAME)
+            .await
+            .unwrap();
+        assert_eq!(delta_table.get_file_uris().unwrap().count(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_optimize_missing_table() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let mut data_folder = DataFolder::open_local(temp_dir.path()).await.unwrap();
+
+        let result = data_folder.optimize(MISSING_TABLE_NAME, None).await;
+
+        assert_eq!(
+            result.unwrap_err().to_string(),
+            format!(
+                "Invalid Argument Error: Table with name '{MISSING_TABLE_NAME}' does not exist."
+            )
+        );
+    }
+
     #[tokio::test]
     async fn test_move_normal_table_to_normal_table() {
         let (_temp_dir, mut source) = create_data_folder_with_normal_table().await;

crates/modelardb_embedded/src/operations/mod.rs

@@ -145,6 +145,15 @@ pub trait Operations: Sync + Send {
         table_name: &str,
         maybe_retention_period_in_seconds: Option<u64>,
     ) -> Result<()>;
+
+    /// Optimize the table with the name in `table_name` by compacting its many small files into
+    /// fewer larger files of approximately `maybe_target_size_in_bytes` bytes. If a target size is
+    /// not given, the default target size of 64 MiB is used.
+    async fn optimize(
+        &mut self,
+        table_name: &str,
+        maybe_target_size_in_bytes: Option<u64>,
+    ) -> Result<()>;
 }
 
 /// Use the time series table metadata in `table_name`, `schema`, `error_bounds`, and `generated_columns`

crates/modelardb_server/src/cluster.rs

@@ -248,6 +248,35 @@ impl Cluster {
         Ok(())
     }
 
+    /// Optimize the tables in `table_names` in the remote data folder and in each peer node. If the
+    /// tables could not be optimized, return [`ModelarDbServerError`].
+    pub(crate) async fn optimize_cluster_tables(
+        &self,
+        table_names: &[String],
+        maybe_target_size_in_bytes: Option<u64>,
+    ) -> Result<()> {
+        // Optimize the tables in the remote data folder.
+        for table_name in table_names {
+            self.remote_data_folder
+                .optimize_table(table_name, maybe_target_size_in_bytes)
+                .await?;
+        }
+
+        // Optimize the tables in each peer node.
+        let optimize_sql = if let Some(target_size_in_bytes) = maybe_target_size_in_bytes {
+            format!(
+                "OPTIMIZE {} TARGET {target_size_in_bytes}",
+                table_names.join(", ")
+            )
+        } else {
+            format!("OPTIMIZE {}", table_names.join(", "))
+        };
+
+        self.cluster_do_get(&optimize_sql).await?;
+
+        Ok(())
+    }
+
     /// For each peer node in the cluster, execute the given `sql` statement with the cluster key
     /// as metadata. If the statement was successfully executed for each node, return [`Ok`],
     /// otherwise return [`ModelarDbServerError`].