1use std::sync::Arc;
21use std::sync::atomic::AtomicU64;
22
23use crate::Result;
24use crate::spec::{DataFileFormat, PartitionKey, TableMetadata, TableProperties};
25use crate::util::location::strip_trailing_slash;
26
27pub trait LocationGenerator: Clone + Send + Sync + 'static {
29 fn generate_location(&self, partition_key: Option<&PartitionKey>, file_name: &str) -> String;
43}
44
45const DEFAULT_DATA_DIR: &str = "/data";
47const HASH_BINARY_STRING_BITS: usize = 20;
49const ENTROPY_DIR_LENGTH: usize = 4;
51const ENTROPY_DIR_DEPTH: usize = 3;
53
54#[derive(Clone, Debug)]
55pub struct DefaultLocationGenerator {
58 data_location: String,
59}
60
61impl DefaultLocationGenerator {
62 pub fn new(table_metadata: &TableMetadata) -> Result<Self> {
69 let table_location = strip_trailing_slash(table_metadata.location());
70 let prop = TableProperties::try_from(table_metadata.properties())?;
71 let data_location = strip_trailing_slash(
72 prop.write_data_location
73 .or(prop.write_folder_storage_location)
74 .unwrap_or(format!("{table_location}{DEFAULT_DATA_DIR}"))
75 .as_ref(),
76 )
77 .to_string();
78
79 Ok(Self { data_location })
80 }
81
82 pub fn with_data_location(data_location: String) -> Self {
88 Self { data_location }
89 }
90}
91
92impl LocationGenerator for DefaultLocationGenerator {
93 fn generate_location(&self, partition_key: Option<&PartitionKey>, file_name: &str) -> String {
94 if PartitionKey::is_effectively_none(partition_key) {
95 format!("{}/{}", self.data_location, file_name)
96 } else {
97 format!(
98 "{}/{}/{}",
99 self.data_location,
100 partition_key.unwrap().to_path(),
101 file_name
102 )
103 }
104 }
105}
106
107#[derive(Clone, Debug)]
121pub struct ObjectStorageLocationGenerator {
122 storage_location: String,
123 context: Option<String>,
125 include_partition_paths: bool,
126}
127
128impl ObjectStorageLocationGenerator {
129 pub fn new(table_metadata: &TableMetadata) -> Result<Self> {
136 let table_location = strip_trailing_slash(table_metadata.location());
137 let prop = TableProperties::try_from(table_metadata.properties())?;
138 let storage_location = strip_trailing_slash(
139 prop.write_data_location
140 .or(prop.write_object_storage_location)
141 .or(prop.write_folder_storage_location)
142 .unwrap_or(format!("{table_location}{DEFAULT_DATA_DIR}"))
143 .as_ref(),
144 )
145 .to_string();
146
147 let context = if storage_location.starts_with(table_location) {
150 None
151 } else {
152 Some(path_context(table_location))
153 };
154
155 let include_partition_paths = prop.write_object_storage_partitioned_paths;
156
157 Ok(Self {
158 storage_location,
159 context,
160 include_partition_paths,
161 })
162 }
163
164 fn new_data_location(&self, name: &str) -> String {
167 let hash = self.compute_hash(name);
168 if let Some(context) = &self.context {
169 format!("{}/{}/{}/{}", self.storage_location, hash, context, name)
170 } else if self.include_partition_paths {
171 format!("{}/{}/{}", self.storage_location, hash, name)
172 } else {
173 format!("{}/{}-{}", self.storage_location, hash, name)
176 }
177 }
178
179 fn compute_hash(&self, name: &str) -> String {
181 let mut bytes = name.as_bytes();
182 let hash_code = murmur3::murmur3_32(&mut bytes, 0).unwrap();
183 let binary = format!("{:032b}", hash_code);
184 let hash = &binary[binary.len() - HASH_BINARY_STRING_BITS..];
185 dirs_from_hash(hash)
186 }
187}
188
189impl LocationGenerator for ObjectStorageLocationGenerator {
190 fn generate_location(&self, partition_key: Option<&PartitionKey>, file_name: &str) -> String {
191 let name =
192 if self.include_partition_paths && !PartitionKey::is_effectively_none(partition_key) {
193 format!("{}/{}", partition_key.unwrap().to_path(), file_name)
194 } else {
195 file_name.to_string()
196 };
197 self.new_data_location(&name)
198 }
199}
200
201fn path_context(table_location: &str) -> String {
204 let mut segments = table_location.rsplit('/').filter(|s| !s.is_empty());
205 let name = segments.next().unwrap_or("");
206 match segments.next() {
207 Some(parent) => format!("{parent}/{name}"),
208 None => name.to_string(),
209 }
210}
211
212fn dirs_from_hash(hash: &str) -> String {
217 let mut result = String::new();
218
219 let mut i = 0;
220 while i < ENTROPY_DIR_DEPTH * ENTROPY_DIR_LENGTH {
221 if i > 0 {
222 result.push('/');
223 }
224 let end = (i + ENTROPY_DIR_LENGTH).min(hash.len());
225 result.push_str(&hash[i..end]);
226 i += ENTROPY_DIR_LENGTH;
227 }
228
229 if hash.len() > ENTROPY_DIR_DEPTH * ENTROPY_DIR_LENGTH {
230 result.push('/');
231 result.push_str(&hash[ENTROPY_DIR_DEPTH * ENTROPY_DIR_LENGTH..]);
232 }
233
234 result
235}
236
237pub trait FileNameGenerator: Clone + Send + Sync + 'static {
239 fn generate_file_name(&self) -> String;
241}
242
243#[derive(Clone, Debug)]
247pub struct DefaultFileNameGenerator {
248 prefix: String,
249 suffix: String,
250 format: String,
251 file_count: Arc<AtomicU64>,
252}
253
254impl DefaultFileNameGenerator {
255 pub fn new(prefix: String, suffix: Option<String>, format: DataFileFormat) -> Self {
257 let suffix = if let Some(suffix) = suffix {
258 format!("-{suffix}")
259 } else {
260 "".to_string()
261 };
262
263 Self {
264 prefix,
265 suffix,
266 format: format.to_string(),
267 file_count: Arc::new(AtomicU64::new(0)),
268 }
269 }
270}
271
272impl FileNameGenerator for DefaultFileNameGenerator {
273 fn generate_file_name(&self) -> String {
274 let file_id = self
275 .file_count
276 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
277 format!(
278 "{}-{:05}{}.{}",
279 self.prefix, file_id, self.suffix, self.format
280 )
281 }
282}
283
284#[cfg(test)]
285pub(crate) mod test {
286 use std::collections::HashMap;
287 use std::sync::Arc;
288
289 use uuid::Uuid;
290
291 use super::LocationGenerator;
292 use crate::spec::{
293 FormatVersion, Literal, NestedField, PartitionKey, PartitionSpec, PrimitiveType, Schema,
294 Struct, StructType, TableMetadata, TableProperties, Transform, Type,
295 };
296 use crate::writer::file_writer::location_generator::{
297 DefaultLocationGenerator, FileNameGenerator, ObjectStorageLocationGenerator,
298 };
299
300 #[test]
301 fn test_default_location_generate() {
302 let mut table_metadata = table_metadata_with("s3://data.db/table", HashMap::new());
303
304 let file_name_generator = super::DefaultFileNameGenerator::new(
305 "part".to_string(),
306 Some("test".to_string()),
307 crate::spec::DataFileFormat::Parquet,
308 );
309
310 let location_generator = DefaultLocationGenerator::new(&table_metadata).unwrap();
312 let location =
313 location_generator.generate_location(None, &file_name_generator.generate_file_name());
314 assert_eq!(location, "s3://data.db/table/data/part-00000-test.parquet");
315
316 table_metadata.properties.insert(
318 TableProperties::PROPERTY_WRITE_FOLDER_STORAGE_LOCATION.to_string(),
319 "s3://data.db/table/data_1".to_string(),
320 );
321 let location_generator = DefaultLocationGenerator::new(&table_metadata).unwrap();
322 let location =
323 location_generator.generate_location(None, &file_name_generator.generate_file_name());
324 assert_eq!(
325 location,
326 "s3://data.db/table/data_1/part-00001-test.parquet"
327 );
328
329 table_metadata.properties.insert(
330 TableProperties::PROPERTY_WRITE_DATA_LOCATION.to_string(),
331 "s3://data.db/table/data_2".to_string(),
332 );
333 let location_generator = DefaultLocationGenerator::new(&table_metadata).unwrap();
334 let location =
335 location_generator.generate_location(None, &file_name_generator.generate_file_name());
336 assert_eq!(
337 location,
338 "s3://data.db/table/data_2/part-00002-test.parquet"
339 );
340
341 table_metadata.properties.insert(
342 TableProperties::PROPERTY_WRITE_DATA_LOCATION.to_string(),
343 "s3://data.db/data_3".to_string(),
345 );
346 let location_generator = DefaultLocationGenerator::new(&table_metadata).unwrap();
347 let location =
348 location_generator.generate_location(None, &file_name_generator.generate_file_name());
349 assert_eq!(location, "s3://data.db/data_3/part-00003-test.parquet");
350 }
351
352 #[test]
353 fn test_location_generate_with_partition() {
354 let schema = Arc::new(
356 Schema::builder()
357 .with_schema_id(1)
358 .with_fields(vec![
359 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
360 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
361 ])
362 .build()
363 .unwrap(),
364 );
365
366 let partition_spec = PartitionSpec::builder(schema.clone())
368 .add_partition_field("id", "id", Transform::Identity)
369 .unwrap()
370 .add_partition_field("name", "name", Transform::Identity)
371 .unwrap()
372 .build()
373 .unwrap();
374
375 let partition_data =
377 Struct::from_iter([Some(Literal::int(42)), Some(Literal::string("alice"))]);
378
379 let partition_key = PartitionKey::new(partition_spec, schema, partition_data);
381
382 let location_gen = DefaultLocationGenerator::with_data_location("/base/path".to_string());
383 let file_name = "data-00000.parquet";
384 let location = location_gen.generate_location(Some(&partition_key), file_name);
385 assert_eq!(location, "/base/path/id=42/name=alice/data-00000.parquet");
386
387 let table_metadata = table_metadata_with("s3://data.db/table", HashMap::new());
389
390 let default_location_gen = DefaultLocationGenerator::new(&table_metadata).unwrap();
392 let location = default_location_gen.generate_location(Some(&partition_key), file_name);
393 assert_eq!(
394 location,
395 "s3://data.db/table/data/id=42/name=alice/data-00000.parquet"
396 );
397 }
398
399 #[test]
400 fn test_location_generate_with_special_characters_partition() {
401 let schema = Arc::new(
402 Schema::builder()
403 .with_schema_id(1)
404 .with_fields(vec![
405 NestedField::required(1, "data#1", Type::Primitive(PrimitiveType::Int)).into(),
406 ])
407 .build()
408 .unwrap(),
409 );
410 let partition_spec = PartitionSpec::builder(schema.clone())
411 .add_partition_field("data#1", "data#1", Transform::Identity)
412 .unwrap()
413 .build()
414 .unwrap();
415 let partition_data = Struct::from_iter([Some(Literal::string("val#1"))]);
416 let partition_key = PartitionKey::new(partition_spec, schema, partition_data);
417
418 let table_metadata = table_metadata_with("s3://data.db/table", HashMap::new());
419 let location_gen = DefaultLocationGenerator::new(&table_metadata).unwrap();
420 let location = location_gen.generate_location(Some(&partition_key), "test.parquet");
421
422 assert_eq!(
423 location,
424 "s3://data.db/table/data/data%231=val%231/test.parquet"
425 );
426 }
427
428 #[test]
430 fn test_object_storage_hash_injection() {
431 let table_metadata = table_metadata_with("s3://data.db/table", HashMap::new());
434 let location_gen = ObjectStorageLocationGenerator::new(&table_metadata).unwrap();
435
436 for (file_name, expected) in [
437 ("a", "s3://data.db/table/data/0101/0110/1001/10110010/a"),
438 ("b", "s3://data.db/table/data/1110/0111/1110/00000011/b"),
439 ("c", "s3://data.db/table/data/0010/1101/0110/01011111/c"),
440 ("d", "s3://data.db/table/data/1001/0001/0100/01110011/d"),
441 ] {
442 assert_eq!(location_gen.generate_location(None, file_name), expected);
443 }
444 }
445
446 #[test]
447 fn test_object_storage_include_partition_paths() {
448 let schema = Arc::new(
451 Schema::builder()
452 .with_schema_id(1)
453 .with_fields(vec![
454 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
455 NestedField::required(2, "a", Type::Primitive(PrimitiveType::String)).into(),
456 ])
457 .build()
458 .unwrap(),
459 );
460 let partition_spec = PartitionSpec::builder(schema.clone())
461 .add_partition_field("id", "id", Transform::Identity)
462 .unwrap()
463 .add_partition_field("a", "a_trunc_3", Transform::Truncate(3))
464 .unwrap()
465 .build()
466 .unwrap();
467 let partition_data =
468 Struct::from_iter([Some(Literal::int(0)), Some(Literal::string("apa"))]);
469 let partition_key = PartitionKey::new(partition_spec, schema, partition_data);
470
471 let table_metadata = table_metadata_with("s3://data.db/table", HashMap::new());
472 let location_gen = ObjectStorageLocationGenerator::new(&table_metadata).unwrap();
473 let location = location_gen.generate_location(Some(&partition_key), "test.parquet");
474
475 assert!(
476 location.starts_with("s3://data.db/table/data/"),
477 "unexpected location: {location}"
478 );
479 assert!(
481 location.ends_with("/id=0/a_trunc_3=apa/test.parquet"),
482 "unexpected location: {location}"
483 );
484 }
485
486 #[test]
488 fn test_object_storage_include_partition_paths_with_special_characters() {
489 let schema = Arc::new(
490 Schema::builder()
491 .with_schema_id(1)
492 .with_fields(vec![
493 NestedField::required(1, "data#1", Type::Primitive(PrimitiveType::Int)).into(),
494 ])
495 .build()
496 .unwrap(),
497 );
498 let partition_spec = PartitionSpec::builder(schema.clone())
499 .add_partition_field("data#1", "data#1", Transform::Identity)
500 .unwrap()
501 .build()
502 .unwrap();
503 let partition_data = Struct::from_iter([Some(Literal::string("val#1"))]);
504 let partition_key = PartitionKey::new(partition_spec, schema, partition_data);
505
506 let table_metadata = table_metadata_with(
507 "s3://data.db/table",
508 HashMap::from([(
509 TableProperties::PROPERTY_WRITE_OBJECT_STORAGE_PARTITIONED_PATHS.to_string(),
510 "true".to_string(),
511 )]),
512 );
513 let location_gen = ObjectStorageLocationGenerator::new(&table_metadata).unwrap();
514 let location = location_gen.generate_location(Some(&partition_key), "test.parquet");
515
516 assert_eq!(
517 location,
518 "s3://data.db/table/data/0000/1011/0110/00001000/data%231=val%231/test.parquet"
519 );
520 }
521
522 #[test]
524 fn test_object_storage_exclude_partition_paths() {
525 let schema = Arc::new(
529 Schema::builder()
530 .with_schema_id(1)
531 .with_fields(vec![
532 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
533 ])
534 .build()
535 .unwrap(),
536 );
537 let partition_spec = PartitionSpec::builder(schema.clone())
538 .add_partition_field("id", "id", Transform::Identity)
539 .unwrap()
540 .build()
541 .unwrap();
542 let partition_data = Struct::from_iter([Some(Literal::int(0))]);
543 let partition_key = PartitionKey::new(partition_spec, schema, partition_data);
544
545 let table_metadata = table_metadata_with(
546 "s3://data.db/table",
547 HashMap::from([(
548 TableProperties::PROPERTY_WRITE_OBJECT_STORAGE_PARTITIONED_PATHS.to_string(),
549 "false".to_string(),
550 )]),
551 );
552 let location_gen = ObjectStorageLocationGenerator::new(&table_metadata).unwrap();
553 let location = location_gen.generate_location(Some(&partition_key), "test.parquet");
554
555 assert_eq!(
556 location,
557 "s3://data.db/table/data/0110/1010/0011/11101000-test.parquet"
558 );
559 }
560
561 #[test]
562 fn test_object_storage_context_when_data_location_outside_table() {
563 let table_metadata = table_metadata_with(
566 "s3://data.db/table",
567 HashMap::from([(
568 TableProperties::PROPERTY_WRITE_DATA_LOCATION.to_string(),
569 "s3://custom-bucket/objects".to_string(),
570 )]),
571 );
572 let location_gen = ObjectStorageLocationGenerator::new(&table_metadata).unwrap();
573 let location = location_gen.generate_location(None, "a");
574
575 assert_eq!(
577 location,
578 "s3://custom-bucket/objects/0101/0110/1001/10110010/data.db/table/a"
579 );
580 }
581
582 #[test]
583 fn test_object_storage_data_location_precedence() {
584 let table_metadata = table_metadata_with(
586 "s3://data.db/table",
587 HashMap::from([
588 (
589 TableProperties::PROPERTY_WRITE_DATA_LOCATION.to_string(),
590 "s3://data.db/table/data_primary".to_string(),
591 ),
592 (
593 TableProperties::PROPERTY_WRITE_FOLDER_STORAGE_LOCATION.to_string(),
594 "s3://data.db/table/data_legacy".to_string(),
595 ),
596 ]),
597 );
598 let location_gen = ObjectStorageLocationGenerator::new(&table_metadata).unwrap();
599 let location = location_gen.generate_location(None, "a");
600
601 assert_eq!(
602 location,
603 "s3://data.db/table/data_primary/0101/0110/1001/10110010/a"
604 );
605 }
606
607 fn table_metadata_with(location: &str, properties: HashMap<String, String>) -> TableMetadata {
609 TableMetadata {
610 format_version: FormatVersion::V2,
611 table_uuid: Uuid::parse_str("fb072c92-a02b-11e9-ae9c-1bb7bc9eca94").unwrap(),
612 location: location.to_string(),
613 last_updated_ms: 1515100955770,
614 last_column_id: 2,
615 schemas: HashMap::new(),
616 current_schema_id: 1,
617 partition_specs: HashMap::new(),
618 default_spec: PartitionSpec::unpartition_spec().into(),
619 default_partition_type: StructType::new(vec![]),
620 last_partition_id: 1000,
621 default_sort_order_id: 0,
622 sort_orders: HashMap::from_iter(vec![]),
623 snapshots: HashMap::default(),
624 current_snapshot_id: None,
625 last_sequence_number: 1,
626 properties,
627 snapshot_log: Vec::new(),
628 metadata_log: vec![],
629 refs: HashMap::new(),
630 statistics: HashMap::new(),
631 partition_statistics: HashMap::new(),
632 encryption_keys: HashMap::new(),
633 next_row_id: 0,
634 }
635 }
636}