Clean DataFolder and extend tests

crates/modelardb_bulkloader/src/main.rs

@@ -38,7 +38,8 @@ use deltalake::{ObjectStore, Path};
 use futures::stream::StreamExt;
 use modelardb_embedded::error::{ModelarDbEmbeddedError, Result};
 use modelardb_embedded::operations::Operations;
-use modelardb_storage::data_folder::{DataFolder, DeltaTableWriter};
+use modelardb_storage::data_folder::DataFolder;
+use modelardb_storage::data_folder::delta_table_writer::DeltaTableWriter;
 use modelardb_types::types::TimeSeriesTableMetadata;
 use sysinfo::System;
 

crates/modelardb_embedded/src/operations/data_folder.rs

@@ -215,15 +215,15 @@ impl Operations for DataFolder {
                 &uncompressed_data,
             )?;
 
-            self.write_compressed_segments_to_time_series_table(table_name, compressed_data)
+            self.write_record_batches(table_name, compressed_data)
                 .await?;
         } else if let Some(normal_table_schema) = self.normal_table_schema(table_name).await {
             // Normal table.
             if !schemas_are_compatible(&uncompressed_data.schema(), &normal_table_schema) {
                 return Err(schema_mismatch_error);
             }
 
-            self.write_record_batches_to_normal_table(table_name, vec![uncompressed_data])
+            self.write_record_batches(table_name, vec![uncompressed_data])
                 .await?;
         } else {
             return Err(ModelarDbEmbeddedError::InvalidArgument(format!(
@@ -282,7 +282,7 @@ impl Operations for DataFolder {
             let record_batches = common::collect(record_batch_stream).await?;
 
             target_data_folder
-                .write_record_batches_to_normal_table(target_table_name, record_batches)
+                .write_record_batches(target_table_name, record_batches)
                 .await?;
 
             Ok(())
@@ -407,7 +407,7 @@ impl Operations for DataFolder {
 
         // Write read data to target_table_name in target.
         target_data_folder
-            .write_compressed_segments_to_time_series_table(target_table_name, record_batches)
+            .write_record_batches(target_table_name, record_batches)
             .await?;
 
         Ok(())
@@ -455,7 +455,7 @@ impl Operations for DataFolder {
             let record_batches: Vec<RecordBatch> = stream.try_collect().await?;
 
             target_data_folder
-                .write_compressed_segments_to_time_series_table(target_table_name, record_batches)
+                .write_record_batches(target_table_name, record_batches)
                 .await?;
         } else if let (Some(source_normal_table_schema), Some(target_normal_table_schema)) = (
             self.normal_table_schema(source_table_name).await,
@@ -474,7 +474,7 @@ impl Operations for DataFolder {
             let record_batches: Vec<RecordBatch> = stream.try_collect().await?;
 
             target_data_folder
-                .write_record_batches_to_normal_table(target_table_name, record_batches)
+                .write_record_batches(target_table_name, record_batches)
                 .await?;
         } else {
             return Err(ModelarDbEmbeddedError::InvalidArgument(format!(

crates/modelardb_server/src/context.rs

@@ -747,10 +747,7 @@ mod tests {
         // Write data to the normal table.
         let local_data_folder = &context.data_folders.local_data_folder;
         local_data_folder
-            .write_record_batches_to_normal_table(
-                NORMAL_TABLE_NAME,
-                vec![table::normal_table_record_batch()],
-            )
+            .write_record_batches(NORMAL_TABLE_NAME, vec![table::normal_table_record_batch()])
             .await
             .unwrap();
 
@@ -818,7 +815,7 @@ mod tests {
         // Write data to the time series table.
         let local_data_folder = &context.data_folders.local_data_folder;
         local_data_folder
-            .write_compressed_segments_to_time_series_table(
+            .write_record_batches(
                 TIME_SERIES_TABLE_NAME,
                 vec![table::compressed_segments_record_batch()],
             )

crates/modelardb_server/src/storage/compressed_data_manager.rs

@@ -87,7 +87,7 @@ impl CompressedDataManager {
         let record_batch_size_in_bytes = record_batch.get_array_memory_size();
 
         self.local_data_folder
-            .write_record_batches_to_normal_table(table_name, vec![record_batch])
+            .write_record_batches(table_name, vec![record_batch])
             .await?;
 
         // Inform the data transfer component about the new data if a remote data folder was
@@ -246,7 +246,7 @@ impl CompressedDataManager {
         let compressed_data_buffer_size_in_bytes = compressed_data_buffer.size_in_bytes;
         let compressed_segments = compressed_data_buffer.record_batches();
         self.local_data_folder
-            .write_compressed_segments_to_time_series_table(table_name, compressed_segments)
+            .write_record_batches(table_name, compressed_segments)
             .await?;
 
         // Inform the data transfer component about the new data if a remote data folder was

crates/modelardb_server/src/storage/data_transfer.rs

@@ -244,19 +244,9 @@ impl DataTransfer {
         debug!("Transferring {current_size_in_bytes} bytes for the table '{table_name}'.",);
 
         // Write the data to the remote Delta Lake.
-        if self
-            .local_data_folder
-            .is_time_series_table(table_name)
-            .await?
-        {
-            self.remote_data_folder
-                .write_compressed_segments_to_time_series_table(table_name, record_batches)
-                .await?;
-        } else {
-            self.remote_data_folder
-                .write_record_batches_to_normal_table(table_name, record_batches)
-                .await?;
-        }
+        self.remote_data_folder
+            .write_record_batches(table_name, record_batches)
+            .await?;
 
         // Delete the data that has been transferred to the remote Delta Lake.
         self.local_data_folder.truncate_table(table_name).await?;
@@ -493,16 +483,13 @@ mod tests {
         for _ in 0..batch_write_count {
             // Write to the normal table.
             local_data_folder
-                .write_record_batches_to_normal_table(
-                    NORMAL_TABLE_NAME,
-                    vec![table::normal_table_record_batch()],
-                )
+                .write_record_batches(NORMAL_TABLE_NAME, vec![table::normal_table_record_batch()])
                 .await
                 .unwrap();
 
             // Write to the time series table.
             local_data_folder
-                .write_compressed_segments_to_time_series_table(
+                .write_record_batches(
                     TIME_SERIES_TABLE_NAME,
                     vec![table::compressed_segments_record_batch()],
                 )