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