Skip to content

Commit c04b7ac

Browse files
pushcoderpiyushh
pushcoder
and
piyushh
authored
refactor(test): remove MockLocationGenerator and its usages (#1645) (#1663)
## Which issue does this PR close? - Closes #1645 - #1645 ## What changes are included in this PR? Refactoring. Keep only one source of location generation. So remove MockLocationGenerator and use DefaultLocationGenerator when table location is available/user given. ## Are these changes tested? Yes. Changed UT of MockLocationGenerator to DefaultLocationGenerator(Using newly added method) Co-authored-by: piyushh <piyush.hurpade@datapelago.com>
1 parent fb94b81 commit c04b7ac

File tree

5 files changed

+73
-70
lines changed

5 files changed

+73
-70
lines changed

‎crates/iceberg/src/writer/base_writer/data_file_writer.rs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -120,16 +120,18 @@ mod test {
120120
};
121121
use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder;
122122
use crate::writer::file_writer::ParquetWriterBuilder;
123-
use crate::writer::file_writer::location_generator::DefaultFileNameGenerator;
124-
use crate::writer::file_writer::location_generator::test::MockLocationGenerator;
123+
use crate::writer::file_writer::location_generator::{
124+
DefaultFileNameGenerator, DefaultLocationGenerator,
125+
};
125126
use crate::writer::{IcebergWriter, IcebergWriterBuilder, RecordBatch};
126127

127128
#[tokio::test]
128129
async fn test_parquet_writer() -> Result<()> {
129130
let temp_dir = TempDir::new().unwrap();
130131
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
131-
let location_gen =
132-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
132+
let location_gen = DefaultLocationGenerator::with_data_location(
133+
temp_dir.path().to_str().unwrap().to_string(),
134+
);
133135
let file_name_gen =
134136
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
135137

@@ -201,8 +203,9 @@ mod test {
201203
async fn test_parquet_writer_with_partition() -> Result<()> {
202204
let temp_dir = TempDir::new().unwrap();
203205
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
204-
let location_gen =
205-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
206+
let location_gen = DefaultLocationGenerator::with_data_location(
207+
temp_dir.path().to_str().unwrap().to_string(),
208+
);
206209
let file_name_gen = DefaultFileNameGenerator::new(
207210
"test_partitioned".to_string(),
208211
None,

‎crates/iceberg/src/writer/base_writer/equality_delete_writer.rs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -198,8 +198,9 @@ mod test {
198198
EqualityDeleteFileWriterBuilder, EqualityDeleteWriterConfig,
199199
};
200200
use crate::writer::file_writer::ParquetWriterBuilder;
201-
use crate::writer::file_writer::location_generator::DefaultFileNameGenerator;
202-
use crate::writer::file_writer::location_generator::test::MockLocationGenerator;
201+
use crate::writer::file_writer::location_generator::{
202+
DefaultFileNameGenerator, DefaultLocationGenerator,
203+
};
203204
use crate::writer::{IcebergWriter, IcebergWriterBuilder};
204205

205206
async fn check_parquet_data_file_with_equality_delete_write(
@@ -282,8 +283,9 @@ mod test {
282283
async fn test_equality_delete_writer() -> Result<(), anyhow::Error> {
283284
let temp_dir = TempDir::new().unwrap();
284285
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
285-
let location_gen =
286-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
286+
let location_gen = DefaultLocationGenerator::with_data_location(
287+
temp_dir.path().to_str().unwrap().to_string(),
288+
);
287289
let file_name_gen =
288290
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
289291

@@ -518,8 +520,9 @@ mod test {
518520
async fn test_equality_delete_with_primitive_type() -> Result<(), anyhow::Error> {
519521
let temp_dir = TempDir::new().unwrap();
520522
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
521-
let location_gen =
522-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
523+
let location_gen = DefaultLocationGenerator::with_data_location(
524+
temp_dir.path().to_str().unwrap().to_string(),
525+
);
523526
let file_name_gen =
524527
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
525528

‎crates/iceberg/src/writer/file_writer/location_generator.rs‎

Lines changed: 13 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,15 @@ impl DefaultLocationGenerator {
6767
};
6868
Ok(Self { data_location })
6969
}
70+
71+
/// Create a new `DefaultLocationGenerator` with a specified data location.
72+
///
73+
/// # Arguments
74+
///
75+
/// * `data_location` - The data location to use for generating file locations.
76+
pub fn with_data_location(data_location: String) -> Self {
77+
Self { data_location }
78+
}
7079
}
7180

7281
impl LocationGenerator for DefaultLocationGenerator {
@@ -144,35 +153,10 @@ pub(crate) mod test {
144153
Struct, StructType, TableMetadata, Transform, Type,
145154
};
146155
use crate::writer::file_writer::location_generator::{
147-
FileNameGenerator, WRITE_DATA_LOCATION, WRITE_FOLDER_STORAGE_LOCATION,
156+
DefaultLocationGenerator, FileNameGenerator, WRITE_DATA_LOCATION,
157+
WRITE_FOLDER_STORAGE_LOCATION,
148158
};
149159

150-
#[derive(Clone)]
151-
pub(crate) struct MockLocationGenerator {
152-
root: String,
153-
}
154-
155-
impl MockLocationGenerator {
156-
pub(crate) fn new(root: String) -> Self {
157-
Self { root }
158-
}
159-
}
160-
161-
impl LocationGenerator for MockLocationGenerator {
162-
fn generate_location(&self, partition: Option<&PartitionKey>, file_name: &str) -> String {
163-
if PartitionKey::is_effectively_none(partition) {
164-
format!("{}/{}", self.root, file_name)
165-
} else {
166-
format!(
167-
"{}/{}/{}",
168-
self.root,
169-
partition.unwrap().to_path(),
170-
file_name
171-
)
172-
}
173-
}
174-
}
175-
176160
#[test]
177161
fn test_default_location_generate() {
178162
let mut table_metadata = TableMetadata {
@@ -283,10 +267,9 @@ pub(crate) mod test {
283267
// Create a partition key
284268
let partition_key = PartitionKey::new(partition_spec, schema, partition_data);
285269

286-
// Test with MockLocationGenerator
287-
let mock_location_gen = MockLocationGenerator::new("/base/path".to_string());
270+
let location_gen = DefaultLocationGenerator::with_data_location("/base/path".to_string());
288271
let file_name = "data-00000.parquet";
289-
let location = mock_location_gen.generate_location(Some(&partition_key), file_name);
272+
let location = location_gen.generate_location(Some(&partition_key), file_name);
290273
assert_eq!(location, "/base/path/id=42/name=alice/data-00000.parquet");
291274

292275
// Create a table metadata for DefaultLocationGenerator

‎crates/iceberg/src/writer/file_writer/parquet_writer.rs‎

Lines changed: 33 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -702,8 +702,9 @@ mod tests {
702702
use crate::arrow::schema_to_arrow_schema;
703703
use crate::io::FileIOBuilder;
704704
use crate::spec::{PrimitiveLiteral, Struct, *};
705-
use crate::writer::file_writer::location_generator::DefaultFileNameGenerator;
706-
use crate::writer::file_writer::location_generator::test::MockLocationGenerator;
705+
use crate::writer::file_writer::location_generator::{
706+
DefaultFileNameGenerator, DefaultLocationGenerator,
707+
};
707708
use crate::writer::tests::check_parquet_data_file;
708709

709710
fn schema_for_all_type() -> Schema {
@@ -862,8 +863,9 @@ mod tests {
862863
async fn test_parquet_writer() -> Result<()> {
863864
let temp_dir = TempDir::new().unwrap();
864865
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
865-
let location_gen =
866-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
866+
let location_gen = DefaultLocationGenerator::with_data_location(
867+
temp_dir.path().to_str().unwrap().to_string(),
868+
);
867869
let file_name_gen =
868870
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
869871

@@ -933,8 +935,9 @@ mod tests {
933935
async fn test_parquet_writer_with_complex_schema() -> Result<()> {
934936
let temp_dir = TempDir::new().unwrap();
935937
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
936-
let location_gen =
937-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
938+
let location_gen = DefaultLocationGenerator::with_data_location(
939+
temp_dir.path().to_str().unwrap().to_string(),
940+
);
938941
let file_name_gen =
939942
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
940943

@@ -1158,8 +1161,9 @@ mod tests {
11581161
async fn test_all_type_for_write() -> Result<()> {
11591162
let temp_dir = TempDir::new().unwrap();
11601163
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
1161-
let loccation_gen =
1162-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
1164+
let loccation_gen = DefaultLocationGenerator::with_data_location(
1165+
temp_dir.path().to_str().unwrap().to_string(),
1166+
);
11631167
let file_name_gen =
11641168
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
11651169

@@ -1401,8 +1405,9 @@ mod tests {
14011405
async fn test_decimal_bound() -> Result<()> {
14021406
let temp_dir = TempDir::new().unwrap();
14031407
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
1404-
let loccation_gen =
1405-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
1408+
let loccation_gen = DefaultLocationGenerator::with_data_location(
1409+
temp_dir.path().to_str().unwrap().to_string(),
1410+
);
14061411
let file_name_gen =
14071412
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
14081413

@@ -1654,8 +1659,9 @@ mod tests {
16541659
async fn test_empty_write() -> Result<()> {
16551660
let temp_dir = TempDir::new().unwrap();
16561661
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
1657-
let location_gen =
1658-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
1662+
let location_gen = DefaultLocationGenerator::with_data_location(
1663+
temp_dir.path().to_str().unwrap().to_string(),
1664+
);
16591665
let file_name_gen =
16601666
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
16611667

@@ -1709,8 +1715,9 @@ mod tests {
17091715
async fn test_nan_val_cnts_primitive_type() -> Result<()> {
17101716
let temp_dir = TempDir::new().unwrap();
17111717
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
1712-
let location_gen =
1713-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
1718+
let location_gen = DefaultLocationGenerator::with_data_location(
1719+
temp_dir.path().to_str().unwrap().to_string(),
1720+
);
17141721
let file_name_gen =
17151722
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
17161723

@@ -1797,8 +1804,9 @@ mod tests {
17971804
async fn test_nan_val_cnts_struct_type() -> Result<()> {
17981805
let temp_dir = TempDir::new().unwrap();
17991806
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
1800-
let location_gen =
1801-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
1807+
let location_gen = DefaultLocationGenerator::with_data_location(
1808+
temp_dir.path().to_str().unwrap().to_string(),
1809+
);
18021810
let file_name_gen =
18031811
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
18041812

@@ -1937,8 +1945,9 @@ mod tests {
19371945
async fn test_nan_val_cnts_list_type() -> Result<()> {
19381946
let temp_dir = TempDir::new().unwrap();
19391947
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
1940-
let location_gen =
1941-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
1948+
let location_gen = DefaultLocationGenerator::with_data_location(
1949+
temp_dir.path().to_str().unwrap().to_string(),
1950+
);
19421951
let file_name_gen =
19431952
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
19441953

@@ -2139,8 +2148,9 @@ mod tests {
21392148
async fn test_nan_val_cnts_map_type() -> Result<()> {
21402149
let temp_dir = TempDir::new().unwrap();
21412150
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
2142-
let location_gen =
2143-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
2151+
let location_gen = DefaultLocationGenerator::with_data_location(
2152+
temp_dir.path().to_str().unwrap().to_string(),
2153+
);
21442154
let file_name_gen =
21452155
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
21462156

@@ -2295,8 +2305,9 @@ mod tests {
22952305
async fn test_write_empty_parquet_file() {
22962306
let temp_dir = TempDir::new().unwrap();
22972307
let file_io = FileIOBuilder::new_fs_io().build().unwrap();
2298-
let location_gen =
2299-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
2308+
let location_gen = DefaultLocationGenerator::with_data_location(
2309+
temp_dir.path().to_str().unwrap().to_string(),
2310+
);
23002311
let file_name_gen =
23012312
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
23022313

‎crates/iceberg/src/writer/file_writer/rolling_writer.rs‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,9 @@ mod tests {
156156
use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, Type};
157157
use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder;
158158
use crate::writer::file_writer::ParquetWriterBuilder;
159-
use crate::writer::file_writer::location_generator::DefaultFileNameGenerator;
160-
use crate::writer::file_writer::location_generator::test::MockLocationGenerator;
159+
use crate::writer::file_writer::location_generator::{
160+
DefaultFileNameGenerator, DefaultLocationGenerator,
161+
};
161162
use crate::writer::tests::check_parquet_data_file;
162163
use crate::writer::{IcebergWriter, IcebergWriterBuilder, RecordBatch};
163164

@@ -188,8 +189,9 @@ mod tests {
188189
async fn test_rolling_writer_basic() -> Result<()> {
189190
let temp_dir = TempDir::new()?;
190191
let file_io = FileIOBuilder::new_fs_io().build()?;
191-
let location_gen =
192-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
192+
let location_gen = DefaultLocationGenerator::with_data_location(
193+
temp_dir.path().to_str().unwrap().to_string(),
194+
);
193195
let file_name_gen =
194196
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
195197

@@ -248,8 +250,9 @@ mod tests {
248250
async fn test_rolling_writer_with_rolling() -> Result<()> {
249251
let temp_dir = TempDir::new()?;
250252
let file_io = FileIOBuilder::new_fs_io().build()?;
251-
let location_gen =
252-
MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
253+
let location_gen = DefaultLocationGenerator::with_data_location(
254+
temp_dir.path().to_str().unwrap().to_string(),
255+
);
253256
let file_name_gen =
254257
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
255258

0 commit comments

Comments
 (0)