diff --git a/crates/modelardb_embedded/bindings/c/modelardb_embedded.h b/crates/modelardb_embedded/bindings/c/modelardb_embedded.h index e4c40b4d4..f8f5f1b56 100644 --- a/crates/modelardb_embedded/bindings/c/modelardb_embedded.h +++ b/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(); diff --git a/crates/modelardb_embedded/bindings/python/modelardb/operations.py b/crates/modelardb_embedded/bindings/python/modelardb/operations.py index f62c3a7a8..b82c87917 100644 --- a/crates/modelardb_embedded/bindings/python/modelardb/operations.py +++ b/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. diff --git a/crates/modelardb_embedded/bindings/python/tests/test_operations.py b/crates/modelardb_embedded/bindings/python/tests/test_operations.py index b5ebb7c43..24fea8315 100644 --- a/crates/modelardb_embedded/bindings/python/tests/test_operations.py +++ b/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()) diff --git a/crates/modelardb_embedded/src/capi.rs b/crates/modelardb_embedded/src/capi.rs index 012d9142f..6e7820659 100644 --- a/crates/modelardb_embedded/src/capi.rs +++ b/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::().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 diff --git a/crates/modelardb_embedded/src/operations/client.rs b/crates/modelardb_embedded/src/operations/client.rs index 6fc70d265..a4d74ed42 100644 --- a/crates/modelardb_embedded/src/operations/client.rs +++ b/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, + ) -> 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(()) + } } diff --git a/crates/modelardb_embedded/src/operations/data_folder.rs b/crates/modelardb_embedded/src/operations/data_folder.rs index e77cf571c..2b45bbfe8 100644 --- a/crates/modelardb_embedded/src/operations/data_folder.rs +++ b/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, + ) -> 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; diff --git a/crates/modelardb_embedded/src/operations/mod.rs b/crates/modelardb_embedded/src/operations/mod.rs index c4e4f467b..2ff33b492 100644 --- a/crates/modelardb_embedded/src/operations/mod.rs +++ b/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, ) -> 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, + ) -> Result<()>; } /// Use the time series table metadata in `table_name`, `schema`, `error_bounds`, and `generated_columns` diff --git a/crates/modelardb_server/src/cluster.rs b/crates/modelardb_server/src/cluster.rs index 6f8eb8e64..8e99d3700 100644 --- a/crates/modelardb_server/src/cluster.rs +++ b/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, + ) -> 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`]. diff --git a/crates/modelardb_server/src/context.rs b/crates/modelardb_server/src/context.rs index 535ee29ce..c94985b5d 100644 --- a/crates/modelardb_server/src/context.rs +++ b/crates/modelardb_server/src/context.rs @@ -401,6 +401,28 @@ impl Context { Ok(()) } + /// Optimize the table with `table_name` if it exists by compacting its small files into larger + /// files of approximately `maybe_target_size_in_bytes` bytes. If a target size is not given, a + /// default target size of 64 MiB is used. If the target size is zero, the table does not exist, + /// or it could not be optimized, [`ModelarDbServerError`] is returned. + pub async fn optimize_table( + &self, + table_name: &str, + maybe_target_size_in_bytes: Option, + ) -> Result<()> { + // Check if the table exists first to provide a consistent error message if it does not. + if self.check_table_does_not_exist(table_name).await.is_ok() { + return Err(table_does_not_exist_error(table_name)); + } + + self.data_folders + .local_data_folder + .optimize_table(table_name, maybe_target_size_in_bytes) + .await?; + + Ok(()) + } + /// Lookup the [`TimeSeriesTableMetadata`] of the time series table with name `table_name` if it /// exists. Specifically, the method returns: /// * [`TimeSeriesTableMetadata`] if a time series table with the name `table_name` exists. @@ -817,25 +839,6 @@ mod tests { assert_eq!(files.count(), 1); } - /// Create a [`Context`] with a normal table named `NORMAL_TABLE_NAME` and write data to it. - async fn create_context_with_normal_table(temp_dir: &TempDir) -> Arc { - let context = create_context(temp_dir).await; - - context - .create_normal_table(NORMAL_TABLE_NAME, &table::normal_table_schema()) - .await - .unwrap(); - - // Write data to the normal table. - let local_data_folder = &context.data_folders.local_data_folder; - local_data_folder - .write_record_batches(NORMAL_TABLE_NAME, vec![table::normal_table_record_batch()]) - .await - .unwrap(); - - context - } - #[tokio::test] async fn test_vacuum_time_series_table() { let temp_dir = tempfile::tempdir().unwrap(); @@ -884,6 +887,120 @@ mod tests { ); } + #[tokio::test] + async fn test_vacuum_missing_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let context = create_context(&temp_dir).await; + + let result = context.vacuum_table("missing_table", None).await; + + assert_eq!( + result.unwrap_err().to_string(), + table_does_not_exist_error("missing_table").to_string() + ); + } + + #[tokio::test] + async fn test_optimize_normal_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let context = create_context_with_normal_table(&temp_dir).await; + + // create_context_with_normal_table() writes one batch. Write three more so there are four + // separate commits, each producing a small file. + let local_data_folder = &context.data_folders.local_data_folder; + for _ in 0..3 { + local_data_folder + .write_record_batches(NORMAL_TABLE_NAME, vec![table::normal_table_record_batch()]) + .await + .unwrap(); + } + + assert_eq!(active_file_count(&context, NORMAL_TABLE_NAME).await, 4); + + context + .optimize_table(NORMAL_TABLE_NAME, None) + .await + .unwrap(); + + // The small files should be compacted into a single active file. + assert_eq!(active_file_count(&context, NORMAL_TABLE_NAME).await, 1); + } + + /// Create a [`Context`] with a normal table named `NORMAL_TABLE_NAME` and write data to it. + async fn create_context_with_normal_table(temp_dir: &TempDir) -> Arc { + let context = create_context(temp_dir).await; + + context + .create_normal_table(NORMAL_TABLE_NAME, &table::normal_table_schema()) + .await + .unwrap(); + + // Write data to the normal table. + let local_data_folder = &context.data_folders.local_data_folder; + local_data_folder + .write_record_batches(NORMAL_TABLE_NAME, vec![table::normal_table_record_batch()]) + .await + .unwrap(); + + context + } + + #[tokio::test] + async fn test_optimize_time_series_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let context = create_context_with_time_series_table(&temp_dir).await; + + // create_context_with_time_series_table() writes one batch. Write three more so there are + // four separate commits, each producing a small file. + let local_data_folder = &context.data_folders.local_data_folder; + for _ in 0..3 { + local_data_folder + .write_record_batches( + TIME_SERIES_TABLE_NAME, + vec![table::compressed_segments_record_batch()], + ) + .await + .unwrap(); + } + + assert_eq!(active_file_count(&context, TIME_SERIES_TABLE_NAME).await, 4); + + context + .optimize_table(TIME_SERIES_TABLE_NAME, None) + .await + .unwrap(); + + // The small files should be compacted into a single active file. + assert_eq!(active_file_count(&context, TIME_SERIES_TABLE_NAME).await, 1); + } + + async fn active_file_count(context: &Context, table_name: &str) -> usize { + let delta_table = context + .data_folders + .local_data_folder + .delta_table(table_name) + .await + .unwrap(); + + delta_table.get_file_uris().unwrap().count() + } + + #[tokio::test] + async fn test_optimize_table_with_zero_target_file_size() { + let temp_dir = tempfile::tempdir().unwrap(); + let context = create_context_with_time_series_table(&temp_dir).await; + + let result = context + .optimize_table(TIME_SERIES_TABLE_NAME, Some(0)) + .await; + + assert_eq!( + result.unwrap_err().to_string(), + "ModelarDB Storage Error: Invalid Argument Error: \ + Optimize target file size must be greater than zero." + ); + } + /// Create a [`Context`] with a time series table named `TIME_SERIES_TABLE_NAME` and write data /// to it. async fn create_context_with_time_series_table(temp_dir: &TempDir) -> Arc { @@ -908,11 +1025,11 @@ mod tests { } #[tokio::test] - async fn test_vacuum_missing_table() { + async fn test_optimize_missing_table() { let temp_dir = tempfile::tempdir().unwrap(); let context = create_context(&temp_dir).await; - let result = context.vacuum_table("missing_table", None).await; + let result = context.optimize_table("missing_table", None).await; assert_eq!( result.unwrap_err().to_string(), diff --git a/crates/modelardb_server/src/remote.rs b/crates/modelardb_server/src/remote.rs index 68733de30..a6a0bdb16 100644 --- a/crates/modelardb_server/src/remote.rs +++ b/crates/modelardb_server/src/remote.rs @@ -469,6 +469,38 @@ impl FlightServiceHandler { Ok(()) } + /// Optimize the tables in `table_names`. If the node is running in a cluster and + /// `optimize_cluster` is `true`, the tables are optimized in the remote data folder and locally + /// in each node in the cluster. If not, the tables are only optimized locally. + async fn optimize_tables( + &self, + table_names: &[String], + maybe_target_size_in_bytes: Option, + optimize_cluster: bool, + ) -> StdResult<(), Status> { + let configuration_manager = self.context.configuration_manager.read().await; + + if optimize_cluster { + if let ClusterMode::MultiNode(cluster) = configuration_manager.cluster_mode() { + cluster + .optimize_cluster_tables(table_names, maybe_target_size_in_bytes) + .await + .map_err(error_to_status_invalid_argument)?; + } else { + return Err(Status::internal("The node is not running in a cluster.")); + } + } + + for table_name in table_names { + self.context + .optimize_table(table_name, maybe_target_size_in_bytes) + .await + .map_err(error_to_status_invalid_argument)?; + } + + Ok(()) + } + /// While there is still more data to receive, ingest the data into the normal table. async fn ingest_into_normal_table( &self, @@ -632,7 +664,7 @@ impl FlightService for FlightServiceHandler { /// Execute a SQL statement provided in UTF-8 and return the schema of the result followed by /// the result itself. Currently, CREATE TABLE, CREATE TIME SERIES TABLE, EXPLAIN, INCLUDE, - /// SELECT, INSERT, TRUNCATE, DROP TABLE, and VACUUM are supported. + /// SELECT, INSERT, TRUNCATE, DROP TABLE, VACUUM, and OPTIMIZE are supported. async fn do_get( &self, request: Request, @@ -723,6 +755,21 @@ impl FlightService for FlightServiceHandler { Ok(empty_record_batch_stream()) } + ModelarDbStatement::Optimize(mut table_names, maybe_target_size_in_bytes, cluster) => { + // Optimize all tables if no table names are provided. + if table_names.is_empty() { + table_names = self + .context + .default_database_schema() + .map_err(error_to_status_internal)? + .table_names(); + }; + + self.optimize_tables(&table_names, maybe_target_size_in_bytes, cluster) + .await?; + + Ok(empty_record_batch_stream()) + } } .map_err(error_to_status_internal)?; diff --git a/crates/modelardb_server/tests/integration_test.rs b/crates/modelardb_server/tests/integration_test.rs index 0ba470a06..e8b998252 100644 --- a/crates/modelardb_server/tests/integration_test.rs +++ b/crates/modelardb_server/tests/integration_test.rs @@ -272,6 +272,22 @@ impl TestContext { self.client.do_get(ticket).await } + /// Optimize a table in the server through the `do_get()` method. + async fn optimize_table( + &mut self, + table_name: &str, + maybe_target_size_in_bytes: Option, + ) -> Result>, Status> { + 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.client.do_get(ticket).await + } + /// Return a [`RecordBatch`] containing a time series with regular or irregular time stamps /// depending on `generate_irregular_timestamps`, generated values with noise depending on /// `multiply_noise_range`, and an optional tag. @@ -801,6 +817,127 @@ async fn test_cannot_vacuum_missing_table() { ); } +#[tokio::test] +async fn test_can_optimize_normal_table() { + let mut test_context = TestContext::new().await; + + let time_series = TestContext::generate_time_series_with_tag(false, None, Some("location")); + ingest_time_series_and_flush_data( + &mut test_context, + slice::from_ref(&time_series), + NORMAL_TABLE_NAME, + TableType::NormalTable, + ) + .await; + + // ingest_time_series_and_flush_data() writes one file. Ingest and flush three more times so + // there are four small files to compact. + for _ in 0..3 { + let flight_data = TestContext::create_flight_data_from_time_series( + NORMAL_TABLE_NAME.to_owned(), + slice::from_ref(&time_series), + ); + + test_context + .send_time_series_to_server(flight_data) + .await + .unwrap(); + + test_context.flush_data_to_disk().await; + } + + let table_path = format!( + "{}/tables/{}", + test_context.temp_dir.path().to_str().unwrap(), + NORMAL_TABLE_NAME + ); + + // Four data files plus the _delta_log folder. + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 5); + + test_context + .optimize_table(NORMAL_TABLE_NAME, None) + .await + .unwrap(); + + // Vacuum to remove the stale files. + test_context + .vacuum_table(NORMAL_TABLE_NAME, Some(0)) + .await + .unwrap(); + + // The active file plus the _delta_log folder should remain. + let files = std::fs::read_dir(&table_path).unwrap(); + assert_eq!(files.count(), 2); +} + +#[tokio::test] +async fn test_can_optimize_time_series_table() { + let mut test_context = TestContext::new().await; + + let time_series = TestContext::generate_time_series_with_tag(false, None, Some("location")); + ingest_time_series_and_flush_data( + &mut test_context, + slice::from_ref(&time_series), + TIME_SERIES_TABLE_NAME, + TableType::TimeSeriesTable, + ) + .await; + + // ingest_time_series_and_flush_data() writes one file per field column partition. ingest and + // flush three more times so each partition has four small files to compact. + for _ in 0..3 { + let flight_data = TestContext::create_flight_data_from_time_series( + TIME_SERIES_TABLE_NAME.to_owned(), + slice::from_ref(&time_series), + ); + + test_context + .send_time_series_to_server(flight_data) + .await + .unwrap(); + + test_context.flush_data_to_disk().await; + } + + let column_path = format!( + "{}/tables/{}/field_column=1", + test_context.temp_dir.path().to_str().unwrap(), + TIME_SERIES_TABLE_NAME + ); + + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 4); + + test_context + .optimize_table(TIME_SERIES_TABLE_NAME, None) + .await + .unwrap(); + + // Vacuum to remove the stale files. + test_context + .vacuum_table(TIME_SERIES_TABLE_NAME, Some(0)) + .await + .unwrap(); + + // The four files in the partition should be compacted into a single active file. + let files = std::fs::read_dir(&column_path).unwrap(); + assert_eq!(files.count(), 1); +} + +#[tokio::test] +async fn test_cannot_optimize_missing_table() { + let mut test_context = TestContext::new().await; + + let response = test_context.optimize_table(NORMAL_TABLE_NAME, None).await; + + assert_eq!( + response.err().unwrap().message(), + format!("Invalid Argument Error: Table with name '{NORMAL_TABLE_NAME}' does not exist.") + ); +} + #[tokio::test] async fn test_can_get_schema() { let mut test_context = TestContext::new().await; diff --git a/crates/modelardb_storage/src/data_folder/mod.rs b/crates/modelardb_storage/src/data_folder/mod.rs index 9bb94eb58..99c8e4492 100644 --- a/crates/modelardb_storage/src/data_folder/mod.rs +++ b/crates/modelardb_storage/src/data_folder/mod.rs @@ -19,6 +19,7 @@ pub mod cluster; pub mod delta_table_writer; use std::collections::{HashMap, HashSet}; +use std::num::NonZeroU64; use std::path::Path as StdPath; use std::sync::Arc; @@ -681,6 +682,32 @@ impl DataFolder { Ok(()) } + /// Optimize the Delta Lake table with `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, a default target size of 64 MiB is used. Compaction only rewrites files smaller + /// than the target, so it is safe to call repeatedly. Note that the small files are only marked + /// as removed and are not deleted from disk until the table is vacuumed. If the target size is + /// zero, the table does not exist, or the files could not be compacted, a + /// [`ModelarDbStorageError`] is returned. + pub async fn optimize_table( + &self, + table_name: &str, + maybe_target_size_in_bytes: Option, + ) -> Result<()> { + let delta_table = self.delta_table(table_name).await?; + + let target_size_in_bytes = maybe_target_size_in_bytes.unwrap_or(64 * 1024 * 1024); + let target_size = NonZeroU64::new(target_size_in_bytes).ok_or_else(|| { + ModelarDbStorageError::InvalidArgument( + "Optimize target file size must be greater than zero.".to_owned(), + ) + })?; + + delta_table.optimize().with_target_size(target_size).await?; + + Ok(()) + } + /// Return a [`DeltaTableWriter`] for writing to the table with `table_name` in the Delta Lake, /// or a [`ModelarDbStorageError`] if a connection to the Delta Lake cannot be established or /// the table does not exist. @@ -1542,6 +1569,139 @@ mod tests { ); } + #[tokio::test] + async fn test_optimize_normal_table() { + let (_temp_dir, data_folder) = create_data_folder_and_create_normal_tables().await; + + // Each write is a separate commit, so four writes produce four small files. + for _ in 0..4 { + data_folder + .write_record_batches("normal_table_1", vec![test::normal_table_record_batch()]) + .await + .unwrap(); + } + + let rows_before = row_count(&data_folder, "normal_table_1").await; + assert_eq!(active_file_count(&data_folder, "normal_table_1").await, 4); + + data_folder + .optimize_table("normal_table_1", None) + .await + .unwrap(); + + // The small files should be compacted into a single active file with no rows lost or + // duplicated. + assert_eq!(active_file_count(&data_folder, "normal_table_1").await, 1); + assert_eq!(row_count(&data_folder, "normal_table_1").await, rows_before); + } + + #[tokio::test] + async fn test_optimize_time_series_table() { + let (_temp_dir, data_folder) = create_data_folder_and_create_time_series_table().await; + + // Each write is a separate commit, so four writes produce four small files. + for _ in 0..4 { + data_folder + .write_record_batches( + TIME_SERIES_TABLE_NAME, + vec![test::compressed_segments_record_batch()], + ) + .await + .unwrap(); + } + + let rows_before = row_count(&data_folder, TIME_SERIES_TABLE_NAME).await; + assert_eq!( + active_file_count(&data_folder, TIME_SERIES_TABLE_NAME).await, + 4 + ); + + data_folder + .optimize_table(TIME_SERIES_TABLE_NAME, None) + .await + .unwrap(); + + // The small files should be compacted into a single active file with no rows lost or + // duplicated. + assert_eq!( + active_file_count(&data_folder, TIME_SERIES_TABLE_NAME).await, + 1 + ); + assert_eq!( + row_count(&data_folder, TIME_SERIES_TABLE_NAME).await, + rows_before + ); + } + + async fn row_count(data_folder: &DataFolder, table_name: &str) -> usize { + let delta_table = data_folder.delta_table(table_name).await.unwrap(); + let (_table, stream) = delta_table.scan_table().await.unwrap(); + + let batches: Vec = stream.try_collect().await.unwrap(); + batches.iter().map(|batch| batch.num_rows()).sum() + } + + #[tokio::test] + async fn test_optimize_table_with_small_target_file_size() { + let (_temp_dir, data_folder) = create_data_folder_and_create_normal_tables().await; + + // Each write is a separate commit, so four writes produce four small files. + for _ in 0..4 { + data_folder + .write_record_batches("normal_table_1", vec![test::normal_table_record_batch()]) + .await + .unwrap(); + } + + let files_before = active_file_count(&data_folder, "normal_table_1").await; + assert_eq!(files_before, 4); + + // A one-byte target is smaller than every existing file, so none of them are candidates for + // compaction, and the files should be left untouched. + data_folder + .optimize_table("normal_table_1", Some(1)) + .await + .unwrap(); + + assert_eq!( + active_file_count(&data_folder, "normal_table_1").await, + files_before + ); + } + + async fn active_file_count(data_folder: &DataFolder, table_name: &str) -> usize { + let delta_table = data_folder.delta_table(table_name).await.unwrap(); + + delta_table.get_file_uris().unwrap().count() + } + + #[tokio::test] + async fn test_optimize_table_with_zero_target_file_size() { + let (_temp_dir, data_folder) = create_data_folder_and_create_normal_tables().await; + + let result = data_folder.optimize_table("normal_table_1", Some(0)).await; + + assert_eq!( + result.unwrap_err().to_string(), + "Invalid Argument Error: Optimize target file size must be greater than zero." + ); + } + + #[tokio::test] + async fn test_optimize_missing_table() { + let (_temp_dir, data_folder) = create_data_folder_and_create_normal_tables().await; + + let result = data_folder.optimize_table("missing_table", None).await; + + assert_eq!( + result.unwrap_err().to_string(), + format!( + "Invalid Argument Error: Delta table cannot be found at '{}'.", + data_folder.location_of_table("missing_table") + ) + ); + } + #[tokio::test] async fn test_write_record_batches_to_normal_table() { let (_temp_dir, data_folder) = create_data_folder_and_create_normal_tables().await; diff --git a/crates/modelardb_storage/src/parser.rs b/crates/modelardb_storage/src/parser.rs index 1841154e1..c3cdaeaa5 100644 --- a/crates/modelardb_storage/src/parser.rs +++ b/crates/modelardb_storage/src/parser.rs @@ -51,8 +51,8 @@ use sqlparser::tokenizer::{Span, Token}; use crate::error::{ModelarDbStorageError, Result}; -/// A top-level statement (CREATE, INSERT, SELECT, TRUNCATE, DROP, VACUUM etc.) that has been -/// tokenized, parsed, and for which semantic checks have verified that it is compatible with +/// A top-level statement (CREATE, INSERT, SELECT, TRUNCATE, DROP, VACUUM, OPTIMIZE etc.) that has +/// been tokenized, parsed, and for which semantic checks have verified that it is compatible with /// ModelarDB. #[derive(Debug)] pub enum ModelarDbStatement { @@ -70,12 +70,14 @@ pub enum ModelarDbStatement { TruncateTable(Vec, bool), /// VACUUM. Vacuum(Vec, Option, bool), + /// OPTIMIZE. + Optimize(Vec, Option, bool), } /// Tokenizes and parses the SQL statement in `sql` and returns its parsed representation in the form /// of a [`ModelarDbStatement`]. Returns a [`ModelarDbStorageError`] if `sql` is empty, contains /// multiple statements, or the statement is unsupported. Currently, CREATE TABLE, CREATE TIME SERIES -/// TABLE, INSERT, EXPLAIN, INCLUDE, SELECT, TRUNCATE, DROP TABLE, and VACUUM are supported. +/// TABLE, INSERT, EXPLAIN, INCLUDE, SELECT, TRUNCATE, DROP TABLE, VACUUM, and OPTIMIZE are supported. pub fn tokenize_and_parse_sql_statement(sql_statement: &str) -> Result { let mut statements = Parser::parse_sql(&ModelarDbDialect::new(), sql_statement)?; @@ -138,10 +140,23 @@ pub fn tokenize_and_parse_sql_statement(sql_statement: &str) -> Result Ok(ModelarDbStatement::Vacuum( - name.value.split_terminator(';').map(|s| s.to_owned()).collect(), - storage_specifier.and_then(|p| p.value.parse::().ok()), + name.value.split_terminator(';').map(|table| table.to_owned()).collect(), + storage_specifier.and_then(|retain| retain.value.parse::().ok()), if_exists, )), + // CreateSecret is used as a substitute for OPTIMIZE since Statement does not have an + // Optimize variant with the required fields. + Statement::CreateSecret { name, storage_specifier, if_not_exists, .. } => { + let table_names = name + .map(|name| name.value.split_terminator(';').map(|table| table.to_owned()).collect()) + .unwrap_or_default(); + + Ok(ModelarDbStatement::Optimize( + table_names, + storage_specifier.and_then(|target| target.value.parse::().ok()), + if_not_exists, + )) + } Statement::Explain { .. } => Ok(ModelarDbStatement::Statement(statement)), Statement::Query(ref boxed_query) => { if let Some(addresses) = extract_include_addresses(boxed_query) { @@ -152,7 +167,8 @@ pub fn tokenize_and_parse_sql_statement(sql_statement: &str) -> Result Ok(ModelarDbStatement::Statement(statement)), _ => Err(ModelarDbStorageError::InvalidArgument( - "Only CREATE, DROP, TRUNCATE, EXPLAIN, INCLUDE, SELECT, INSERT, and VACUUM are supported." + "Only CREATE, DROP, TRUNCATE, EXPLAIN, INCLUDE, SELECT, INSERT, VACUUM, and OPTIMIZE \ + are supported." .to_owned(), )), } @@ -179,7 +195,8 @@ pub fn tokenize_and_parse_sql_expression( /// SQL dialect that extends `sqlparsers's` [`GenericDialect`] with support for parsing CREATE TIME /// SERIES TABLE table_name DDL statements, INCLUDE 'address'\[, 'address'\]+ DQL statements, -/// VACUUM \[CLUSTER\] \[table_name\[, table_name\]+\] \[RETAIN num_seconds\] statements, and +/// VACUUM \[CLUSTER\] \[table_name\[, table_name\]+\] \[RETAIN num_seconds\] statements, +/// OPTIMIZE \[CLUSTER\] \[table_name\[, table_name\]+\] \[TARGET num_bytes\] statements, and /// TRUNCATE \[CLUSTER\] table_name\[, table_name\]+ statements. #[derive(Debug)] struct ModelarDbDialect { @@ -561,6 +578,78 @@ impl ModelarDbDialect { }) } + /// Return [`true`] if the token stream starts with OPTIMIZE, otherwise [`false`] is returned. + /// The method does not consume tokens. + fn next_token_is_optimize(&self, parser: &Parser) -> bool { + // OPTIMIZE. + if let Token::Word(word) = parser.peek_nth_token(0).token { + word.keyword == Keyword::OPTIMIZE + } else { + false + } + } + + /// Parse OPTIMIZE \[CLUSTER\] \[table_name\[, table_name\]+\] \[TARGET num_bytes\] to a + /// [`Statement::CreateSecret`] with the table names in the `name` field, the cluster flag in + /// the `if_not_exists` field, and the optional target file size in the `storage_specifier` + /// field. Note that [`Statement::CreateSecret`] is used since [`Statement`] does not have an + /// `Optimize` variant with the required fields. A [`ParserError`] is returned if OPTIMIZE is + /// not the first word, the table names cannot be extracted, or the target file size is not a + /// valid positive integer. + fn parse_optimize(&self, parser: &mut Parser) -> StdResult { + // OPTIMIZE. + parser.expect_keyword(Keyword::OPTIMIZE)?; + + // If the next token is CLUSTER, consume it and set the flag to optimize the entire cluster. + let optimize_cluster = if let Token::Word(word) = parser.peek_nth_token(0).token + && word.keyword == Keyword::CLUSTER + { + parser.expect_keyword(Keyword::CLUSTER)?; + true + } else { + false + }; + + // If the next token is a word that is not TARGET, attempt to parse table names. + let table_names = if let Token::Word(word) = parser.peek_nth_token(0).token + && word.keyword != Keyword::TARGET + { + self.parse_table_names(parser)? + } else { + vec![] + }; + + // If the next token is TARGET, attempt to parse the target file size in bytes. + let maybe_target_size_in_bytes = if let Token::Word(word) = parser.peek_nth_token(0).token + && word.keyword == Keyword::TARGET + { + parser.expect_keyword(Keyword::TARGET)?; + let target_size_in_bytes = self.parse_unsigned_literal_u64(parser)?; + + if target_size_in_bytes == 0 { + return Err(ParserError::ParserError( + "Target file size must be a positive integer.".to_owned(), + )); + } + + Some(target_size_in_bytes) + } else { + None + }; + + // Return Statement::CreateSecret as a substitute for Optimize. + Ok(Statement::CreateSecret { + if_not_exists: optimize_cluster, + name: Some(Ident::new(table_names.join(";"))), + storage_specifier: maybe_target_size_in_bytes + .map(|target| Ident::new(target.to_string())), + secret_type: Ident::new(""), + options: vec![], + or_replace: false, + temporary: None, + }) + } + /// Return its value as a [`u64`] if the next [`Token`] is a [`Token::Number`], otherwise a /// [`ParserError`] is returned. fn parse_unsigned_literal_u64(&self, parser: &mut Parser) -> StdResult { @@ -688,6 +777,8 @@ impl Dialect for ModelarDbDialect { /// attempt to parse the token stream as an INCLUDE 'address'\[, 'address'\]+ DQL statement. /// If not, check if the next token is VACUUM, if so, attempt to parse the token stream as a /// VACUUM \[CLUSTER\] \[table_name\[, table_name\]+\] \[RETAIN num_seconds\] statement. If not, + /// check if the next token is OPTIMIZE, if so, attempt to parse the token stream as an + /// OPTIMIZE \[CLUSTER\] \[table_name\[, table_name\]+\] \[TARGET num_bytes\] statement. If not, /// check if the next token is TRUNCATE, if so, attempt to parse the token stream as a /// TRUNCATE \[CLUSTER\] table_name\[, table_name\]+ statement. If all checks fail, [`None`] is /// returned so [`sqlparser`] uses its parsing methods for all other statements. If parsing @@ -699,6 +790,8 @@ impl Dialect for ModelarDbDialect { Some(self.parse_include_query(parser)) } else if self.next_token_is_vacuum(parser) { Some(self.parse_vacuum(parser)) + } else if self.next_token_is_optimize(parser) { + Some(self.parse_optimize(parser)) } else if self.next_token_is_truncate(parser) { Some(self.parse_truncate(parser)) } else { @@ -2180,6 +2273,307 @@ mod tests { ); } + #[test] + fn test_tokenize_and_parse_optimize_all_tables() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names("OPTIMIZE"); + + assert!(table_names.is_empty()); + assert!(maybe_target_size_in_bytes.is_none()); + assert!(!cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_single_table() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names("OPTIMIZE table_name"); + + assert_eq!(table_names, vec!["table_name".to_owned()]); + assert!(maybe_target_size_in_bytes.is_none()); + assert!(!cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_multiple_tables() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names("OPTIMIZE table_name_1, table_name_2"); + + assert_eq!( + table_names, + vec!["table_name_1".to_owned(), "table_name_2".to_owned()] + ); + assert!(maybe_target_size_in_bytes.is_none()); + assert!(!cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_with_target_size() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names("OPTIMIZE TARGET 1024"); + + assert!(table_names.is_empty()); + assert_eq!(maybe_target_size_in_bytes, Some(1024)); + assert!(!cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_multiple_tables_with_target_size() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names( + "OPTIMIZE table_name_1, table_name_2 TARGET 1024", + ); + + assert_eq!( + table_names, + vec!["table_name_1".to_owned(), "table_name_2".to_owned()] + ); + assert_eq!(maybe_target_size_in_bytes, Some(1024)); + assert!(!cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_cluster() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names("OPTIMIZE CLUSTER"); + + assert!(table_names.is_empty()); + assert!(maybe_target_size_in_bytes.is_none()); + assert!(cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_cluster_with_multiple_tables() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names("OPTIMIZE CLUSTER table_name_1, table_name_2"); + + assert_eq!( + table_names, + vec!["table_name_1".to_owned(), "table_name_2".to_owned()] + ); + assert!(maybe_target_size_in_bytes.is_none()); + assert!(cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_cluster_with_target_size() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names("OPTIMIZE CLUSTER TARGET 1024"); + + assert!(table_names.is_empty()); + assert_eq!(maybe_target_size_in_bytes, Some(1024)); + assert!(cluster) + } + + #[test] + fn test_tokenize_and_parse_optimize_cluster_with_multiple_tables_and_target_size() { + let (table_names, maybe_target_size_in_bytes, cluster) = + parse_optimize_and_extract_table_names( + "OPTIMIZE CLUSTER table_name_1, table_name_2 TARGET 1024", + ); + + assert_eq!( + table_names, + vec!["table_name_1".to_owned(), "table_name_2".to_owned()] + ); + assert_eq!(maybe_target_size_in_bytes, Some(1024)); + assert!(cluster) + } + + fn parse_optimize_and_extract_table_names( + sql_statement: &str, + ) -> (Vec, Option, bool) { + let modelardb_statement = tokenize_and_parse_sql_statement(sql_statement).unwrap(); + + match modelardb_statement { + ModelarDbStatement::Optimize(table_names, maybe_target_size_in_bytes, cluster) => { + (table_names, maybe_target_size_in_bytes, cluster) + } + _ => panic!("Expected ModelarDbStatement::Optimize."), + } + } + + #[test] + fn test_tokenize_and_parse_optimize_trailing_comma() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE table_name,"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: word, found: EOF" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_leading_comma() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE ,table_name"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: , at Line: 1, Column: 10" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_only_comma() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE,"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: , at Line: 1, Column: 9" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_quoted_table_name() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE 'table_name'"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: 'table_name' at Line: 1, Column: 10" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_without_number() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: literal integer, found: EOF" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_number_without_target() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE 1024"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: 1024 at Line: 1, Column: 10" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_with_float() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET 1024.5"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Failed to parse '1024.5' into a u64 due to: invalid digit found in string" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_with_non_numeric() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET thousand"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: literal integer, found: thousand at Line: 1, Column: 17" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_with_negative() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET -1024"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: literal integer, found: - at Line: 1, Column: 17" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_with_zero() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET 0"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Target file size must be a positive integer." + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_twice() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET 1024 TARGET 1024"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: TARGET at Line: 1, Column: 22" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_multiple_tables_target_without_number() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE table_1, table_2 TARGET"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: literal integer, found: EOF" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_tables_and_target_mixed() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE table_1, TARGET 1024, table_2"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: 1024 at Line: 1, Column: 26" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_first() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET 1024 table_1, table_2"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: table_1 at Line: 1, Column: 22" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_cluster_last() { + let result = + tokenize_and_parse_sql_statement("OPTIMIZE table_1, table_2 TARGET 1024 CLUSTER"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: CLUSTER at Line: 1, Column: 39" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_cluster_after_tables_before_target() { + let result = + tokenize_and_parse_sql_statement("OPTIMIZE table_1, table_2 CLUSTER TARGET 1024"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: CLUSTER at Line: 1, Column: 27" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_cluster_trailing_comma() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE CLUSTER, table_1"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: end of statement, found: , at Line: 1, Column: 17" + ); + } + + #[test] + fn test_tokenize_and_parse_optimize_target_and_cluster_mixed() { + let result = tokenize_and_parse_sql_statement("OPTIMIZE TARGET CLUSTER 1024"); + + assert_eq!( + result.unwrap_err().to_string(), + "Parser Error: sql parser error: Expected: literal integer, found: CLUSTER at Line: 1, Column: 17" + ); + } + #[test] fn test_tokenize_and_parse_truncate_single_table() { let (table_names, cluster) = parse_truncate_and_extract_table_names("TRUNCATE table_name");