1use std::{collections::HashMap, env, hash::RandomState, sync::Arc};
5
6#[cfg(test)]
7use arrow_array::cast::AsArray;
8use arrow_array::{ArrayRef, UInt8Array};
9use arrow_schema::DataType;
10use hyperloglogplus::{HyperLogLog, HyperLogLogPlus};
11
12use crate::{
13 array_encoding::{
14 logical::{
15 blob::BlobFieldEncoder, list::ListFieldEncoder, primitive::PrimitiveFieldEncoder,
16 },
17 physical::{
18 basic::BasicEncoder,
19 binary::BinaryEncoder,
20 dictionary::{AlreadyDictionaryEncoder, DictionaryEncoder},
21 fixed_size_list::FslEncoder,
22 fsst::FsstArrayEncoder,
23 packed_struct::PackedStructEncoder,
24 },
25 },
26 constants::{
27 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, PACKED_STRUCT_LEGACY_META_KEY,
28 PACKED_STRUCT_META_KEY,
29 },
30 encoder::{
31 ArrayEncoder, ArrayEncodingStrategy, ColumnIndexSequence, FieldEncoder,
32 FieldEncodingContext, FieldEncodingStrategy,
33 },
34 encodings::{
35 logical::r#struct::StructFieldEncoder,
36 physical::{
37 block::{CompressionConfig, CompressionScheme},
38 value::ValueEncoder,
39 },
40 },
41};
42
43use lance_arrow::BLOB_META_KEY;
44use lance_core::datatypes::{BLOB_DESC_FIELD, Field};
45use lance_core::{Error, Result};
46
47#[derive(Debug)]
49pub struct ArrayFieldEncodingStrategy {
50 array_encoding_strategy: Arc<dyn ArrayEncodingStrategy>,
51}
52
53impl ArrayFieldEncodingStrategy {
54 pub fn new() -> Self {
62 Self {
63 array_encoding_strategy: Arc::new(ArrayStrategy),
64 }
65 }
66
67 fn is_primitive_type(data_type: &DataType) -> bool {
68 matches!(
69 data_type,
70 DataType::Boolean
71 | DataType::Date32
72 | DataType::Date64
73 | DataType::Decimal128(_, _)
74 | DataType::Decimal256(_, _)
75 | DataType::Duration(_)
76 | DataType::Float16
77 | DataType::Float32
78 | DataType::Float64
79 | DataType::Int16
80 | DataType::Int32
81 | DataType::Int64
82 | DataType::Int8
83 | DataType::Interval(_)
84 | DataType::Null
85 | DataType::Time32(_)
86 | DataType::Time64(_)
87 | DataType::Timestamp(_, _)
88 | DataType::UInt16
89 | DataType::UInt32
90 | DataType::UInt64
91 | DataType::UInt8
92 | DataType::FixedSizeBinary(_)
93 | DataType::FixedSizeList(_, _)
94 | DataType::Binary
95 | DataType::LargeBinary
96 | DataType::Utf8
97 | DataType::LargeUtf8,
98 )
99 }
100}
101
102impl Default for ArrayFieldEncodingStrategy {
103 fn default() -> Self {
104 Self::new()
105 }
106}
107
108impl FieldEncodingStrategy for ArrayFieldEncodingStrategy {
109 fn create_field_encoder(
110 &self,
111 field: &Field,
112 column_index: &mut ColumnIndexSequence,
113 context: &FieldEncodingContext<'_>,
114 ) -> Result<Box<dyn FieldEncoder>> {
115 let options = context.options;
116 let data_type = field.data_type();
117 if Self::is_primitive_type(&data_type) {
118 let column_index = column_index.next_column_index(field.id as u32);
119 if field.metadata.contains_key(BLOB_META_KEY) {
120 let mut packed_meta = HashMap::new();
121 packed_meta.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string());
122 let desc_field =
123 Field::try_from(BLOB_DESC_FIELD.clone().with_metadata(packed_meta)).unwrap();
124 let desc_encoder = Box::new(PrimitiveFieldEncoder::try_new(
125 options,
126 self.array_encoding_strategy.clone(),
127 column_index,
128 desc_field,
129 )?);
130 Ok(Box::new(BlobFieldEncoder::new(desc_encoder)))
131 } else {
132 Ok(Box::new(PrimitiveFieldEncoder::try_new(
133 options,
134 self.array_encoding_strategy.clone(),
135 column_index,
136 field.clone(),
137 )?))
138 }
139 } else {
140 match data_type {
141 DataType::List(_child) | DataType::LargeList(_child) => {
142 let list_idx = column_index.next_column_index(field.id as u32);
143 let inner_encoding = context.strategy.create_field_encoder(
144 &field.children[0],
145 column_index,
146 context,
147 )?;
148 let offsets_encoder =
149 Arc::new(BasicEncoder::new(Box::new(ValueEncoder::default())));
150 Ok(Box::new(ListFieldEncoder::new(
151 inner_encoding,
152 offsets_encoder,
153 options.cache_bytes_per_column,
154 options.keep_original_array,
155 list_idx,
156 )))
157 }
158 DataType::Struct(_) => {
159 let field_metadata = &field.metadata;
160 if field_metadata
161 .get(PACKED_STRUCT_LEGACY_META_KEY)
162 .map(|v| v == "true")
163 .unwrap_or(field_metadata.contains_key(PACKED_STRUCT_META_KEY))
164 {
165 Ok(Box::new(PrimitiveFieldEncoder::try_new(
166 options,
167 self.array_encoding_strategy.clone(),
168 column_index.next_column_index(field.id as u32),
169 field.clone(),
170 )?))
171 } else {
172 let header_idx = column_index.next_column_index(field.id as u32);
173 let children_encoders = field
174 .children
175 .iter()
176 .map(|field| {
177 context
178 .strategy
179 .create_field_encoder(field, column_index, context)
180 })
181 .collect::<Result<Vec<_>>>()?;
182 Ok(Box::new(StructFieldEncoder::new(
183 children_encoders,
184 header_idx,
185 )))
186 }
187 }
188 DataType::Dictionary(_, value_type) => {
189 if Self::is_primitive_type(&value_type) {
191 Ok(Box::new(PrimitiveFieldEncoder::try_new(
192 options,
193 self.array_encoding_strategy.clone(),
194 column_index.next_column_index(field.id as u32),
195 field.clone(),
196 )?))
197 } else {
198 Err(Error::not_supported_source(format!("cannot encode a dictionary column whose value type is a logical type ({})", value_type).into()))
204 }
205 }
206 _ => Err(Error::not_supported_source(
207 format!(
208 "Lance v2.0 has no field encoding for '{}' with data type {}",
209 field.name,
210 field.data_type()
211 )
212 .into(),
213 )),
214 }
215 }
216 }
217}
218
219#[derive(Debug)]
221struct ArrayStrategy;
222
223impl ArrayStrategy {
224 fn get_field_compression(field_meta: &HashMap<String, String>) -> Option<CompressionConfig> {
225 let compression = field_meta.get(COMPRESSION_META_KEY)?;
226 let compression_scheme = compression.parse::<CompressionScheme>();
227 match compression_scheme {
228 Ok(compression_scheme) => Some(CompressionConfig::new(
229 compression_scheme,
230 field_meta
231 .get(COMPRESSION_LEVEL_META_KEY)
232 .and_then(|level| level.parse().ok()),
233 )),
234 Err(_) => None,
235 }
236 }
237
238 fn default_binary_encoder(
239 arrays: &[ArrayRef],
240 field_meta: Option<&HashMap<String, String>>,
241 data_size: u64,
242 ) -> Result<Box<dyn ArrayEncoder>> {
243 let bin_indices_encoder =
244 Self::choose_array_encoder(arrays, &DataType::UInt64, data_size, false, None)?;
245
246 if let Some(compression) = field_meta.and_then(Self::get_field_compression) {
247 if compression.scheme() == CompressionScheme::Fsst {
248 let raw_encoder = Box::new(BinaryEncoder::try_new(bin_indices_encoder, None)?);
250 Ok(Box::new(FsstArrayEncoder::new(raw_encoder)))
251 } else {
252 Ok(Box::new(BinaryEncoder::try_new(
254 bin_indices_encoder,
255 Some(compression),
256 )?))
257 }
258 } else {
259 Ok(Box::new(BinaryEncoder::try_new(bin_indices_encoder, None)?))
260 }
261 }
262
263 fn choose_array_encoder(
264 arrays: &[ArrayRef],
265 data_type: &DataType,
266 data_size: u64,
267 use_dict_encoding: bool,
268 field_meta: Option<&HashMap<String, String>>,
269 ) -> Result<Box<dyn ArrayEncoder>> {
270 match data_type {
271 DataType::FixedSizeList(inner, dimension) => {
272 Ok(Box::new(BasicEncoder::new(Box::new(FslEncoder::new(
273 Self::choose_array_encoder(
274 arrays,
275 inner.data_type(),
276 data_size,
277 use_dict_encoding,
278 None,
279 )?,
280 *dimension as u32,
281 )))))
282 }
283 DataType::Dictionary(key_type, value_type) => {
284 let key_encoder =
285 Self::choose_array_encoder(arrays, key_type, data_size, false, None)?;
286 let value_encoder =
287 Self::choose_array_encoder(arrays, value_type, data_size, false, None)?;
288
289 Ok(Box::new(AlreadyDictionaryEncoder::new(
290 key_encoder,
291 value_encoder,
292 )))
293 }
294 DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
295 if use_dict_encoding {
296 let dict_indices_encoder = Self::choose_array_encoder(
297 &[Arc::new(UInt8Array::from_iter_values(0_u8..255_u8))],
302 &DataType::UInt8,
303 data_size,
304 false,
305 None,
306 )?;
307 let dict_items_encoder = Self::choose_array_encoder(
308 arrays,
309 &DataType::Utf8,
310 data_size,
311 false,
312 None,
313 )?;
314
315 Ok(Box::new(DictionaryEncoder::new(
316 dict_indices_encoder,
317 dict_items_encoder,
318 )))
319 } else {
320 Self::default_binary_encoder(arrays, field_meta, data_size)
321 }
322 }
323 DataType::Struct(fields) => {
324 let num_fields = fields.len();
325 let mut inner_encoders = Vec::new();
326
327 for i in 0..num_fields {
328 let inner_datatype = fields[i].data_type();
329 let inner_encoder = Self::choose_array_encoder(
330 arrays,
331 inner_datatype,
332 data_size,
333 use_dict_encoding,
334 None,
335 )?;
336 inner_encoders.push(inner_encoder);
337 }
338
339 Ok(Box::new(PackedStructEncoder::new(inner_encoders)))
340 }
341 DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => Ok(
342 Box::new(BasicEncoder::new(Box::new(ValueEncoder::default()))),
343 ),
344
345 DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Box::new(
349 BasicEncoder::new(Box::new(ValueEncoder::default())),
350 )),
351 _ => Ok(Box::new(BasicEncoder::new(Box::new(
352 ValueEncoder::default(),
353 )))),
354 }
355 }
356}
357
358fn get_dict_encoding_threshold() -> u64 {
359 env::var("LANCE_DICT_ENCODING_THRESHOLD")
360 .ok()
361 .and_then(|val| val.parse().ok())
362 .unwrap_or(100)
363}
364
365fn check_dict_encoding(arrays: &[ArrayRef], threshold: u64) -> bool {
373 let num_total_rows = arrays.iter().map(|arr| arr.len()).sum::<usize>();
374 if num_total_rows < threshold as usize {
375 return false;
376 }
377 const PRECISION: u8 = 12;
378
379 let mut hll: HyperLogLogPlus<String, RandomState> =
380 HyperLogLogPlus::new(PRECISION, RandomState::new()).unwrap();
381
382 for arr in arrays {
383 let string_array = arrow_array::cast::as_string_array(arr);
384 for value in string_array.iter().flatten() {
385 hll.insert(value);
386 let estimated_cardinality = hll.count() as u64;
387 if estimated_cardinality >= threshold {
388 return false;
389 }
390 }
391 }
392
393 true
394}
395
396#[cfg(test)]
397fn check_fixed_size_encoding(arrays: &[ArrayRef]) -> Option<u64> {
398 if arrays.is_empty() {
399 return None;
400 }
401
402 if !arrays.iter().all(|arr| {
404 if let Some(arr) = arr.as_string_opt::<i32>() {
405 arr.iter().flatten().all(|s| !s.is_empty())
406 } else if let Some(arr) = arr.as_binary_opt::<i32>() {
407 arr.iter().flatten().all(|s| !s.is_empty())
408 } else if let Some(arr) = arr.as_string_opt::<i64>() {
409 arr.iter().flatten().all(|s| !s.is_empty())
410 } else if let Some(arr) = arr.as_binary_opt::<i64>() {
411 arr.iter().flatten().all(|s| !s.is_empty())
412 } else {
413 panic!("wrong dtype");
414 }
415 }) {
416 return None;
417 }
418
419 let lengths = arrays
420 .iter()
421 .flat_map(|arr| {
422 if let Some(arr) = arr.as_string_opt::<i32>() {
423 let offsets = arr.offsets().inner();
424 offsets
425 .windows(2)
426 .map(|w| (w[1] - w[0]) as u64)
427 .collect::<Vec<_>>()
428 } else if let Some(arr) = arr.as_binary_opt::<i32>() {
429 let offsets = arr.offsets().inner();
430 offsets
431 .windows(2)
432 .map(|w| (w[1] - w[0]) as u64)
433 .collect::<Vec<_>>()
434 } else if let Some(arr) = arr.as_string_opt::<i64>() {
435 let offsets = arr.offsets().inner();
436 offsets
437 .windows(2)
438 .map(|w| (w[1] - w[0]) as u64)
439 .collect::<Vec<_>>()
440 } else if let Some(arr) = arr.as_binary_opt::<i64>() {
441 let offsets = arr.offsets().inner();
442 offsets
443 .windows(2)
444 .map(|w| (w[1] - w[0]) as u64)
445 .collect::<Vec<_>>()
446 } else {
447 panic!("wrong dtype");
448 }
449 })
450 .collect::<Vec<_>>();
451
452 let first_non_zero = lengths.iter().position(|&x| x != 0);
454 if let Some(first_non_zero) = first_non_zero {
455 if !lengths
457 .iter()
458 .all(|&x| x == 0 || x == lengths[first_non_zero])
459 {
460 return None;
461 }
462
463 Some(lengths[first_non_zero])
465 } else {
466 None
467 }
468}
469
470impl ArrayEncodingStrategy for ArrayStrategy {
471 fn create_array_encoder(
472 &self,
473 arrays: &[ArrayRef],
474 field: &Field,
475 ) -> Result<Box<dyn ArrayEncoder>> {
476 let data_size = arrays
477 .iter()
478 .map(|arr| arr.get_buffer_memory_size() as u64)
479 .sum::<u64>();
480 let data_type = arrays[0].data_type();
481
482 let use_dict_encoding = data_type == &DataType::Utf8
483 && check_dict_encoding(arrays, get_dict_encoding_threshold());
484
485 Self::choose_array_encoder(
486 arrays,
487 data_type,
488 data_size,
489 use_dict_encoding,
490 Some(&field.metadata),
491 )
492 }
493}
494
495#[cfg(test)]
496mod tests {
497 use super::{
498 ArrayEncodingStrategy, ArrayFieldEncodingStrategy, ArrayStrategy, check_dict_encoding,
499 check_fixed_size_encoding,
500 };
501 use crate::constants::{COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY};
502 use crate::encoder::{BatchEncoder, EncodingOptions};
503 use arrow_array::{ArrayRef, StringArray};
504 use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema};
505 use lance_core::{Error, datatypes::Schema};
506 use std::collections::HashMap;
507 use std::sync::Arc;
508
509 #[test]
510 fn test_unsupported_field_type_returns_error() {
511 let entries = Field::new(
512 "entries",
513 DataType::Struct(Fields::from(vec![
514 Field::new("key", DataType::Utf8, false),
515 Field::new("value", DataType::Int32, true),
516 ])),
517 false,
518 );
519 let arrow_schema = ArrowSchema::new(vec![Field::new(
520 "attributes",
521 DataType::Map(Arc::new(entries), false),
522 true,
523 )]);
524 let schema = Schema::try_from(&arrow_schema).unwrap();
525
526 let error = BatchEncoder::try_new(
527 &schema,
528 &ArrayFieldEncodingStrategy::new(),
529 &EncodingOptions::default(),
530 )
531 .err()
532 .unwrap();
533
534 assert!(matches!(error, Error::NotSupported { .. }));
535 assert!(error.to_string().contains("attributes"));
536 assert!(error.to_string().contains("Map"));
537 }
538
539 fn is_dict_encoding_applicable(arr: Vec<Option<&str>>, threshold: u64) -> bool {
540 let arr = StringArray::from(arr);
541 let arr = Arc::new(arr) as ArrayRef;
542 check_dict_encoding(&[arr], threshold)
543 }
544
545 #[test]
546 fn test_dict_encoding_should_be_applied_if_cardinality_less_than_threshold() {
547 assert!(is_dict_encoding_applicable(
548 vec![Some("a"), Some("b"), Some("a"), Some("b")],
549 3,
550 ));
551 }
552
553 #[test]
554 fn test_dict_encoding_should_not_be_applied_if_cardinality_larger_than_threshold() {
555 assert!(!is_dict_encoding_applicable(
556 vec![Some("a"), Some("b"), Some("c"), Some("d")],
557 3,
558 ));
559 }
560
561 #[test]
562 fn test_dict_encoding_should_not_be_applied_if_cardinality_equal_to_threshold() {
563 assert!(!is_dict_encoding_applicable(
564 vec![Some("a"), Some("b"), Some("c"), Some("a")],
565 3,
566 ));
567 }
568
569 #[test]
570 fn test_dict_encoding_should_not_be_applied_for_empty_arrays() {
571 assert!(!is_dict_encoding_applicable(vec![], 3));
572 }
573
574 #[test]
575 fn test_dict_encoding_should_not_be_applied_for_smaller_than_threshold_arrays() {
576 assert!(!is_dict_encoding_applicable(vec![Some("a"), Some("a")], 3));
577 }
578
579 fn is_fixed_size_encoding_applicable(arrays: Vec<Vec<Option<&str>>>) -> bool {
580 let mut final_arrays = Vec::new();
581 for arr in arrays {
582 let arr = StringArray::from(arr);
583 let arr = Arc::new(arr) as ArrayRef;
584 final_arrays.push(arr);
585 }
586
587 check_fixed_size_encoding(&final_arrays).is_some()
588 }
589
590 #[test]
591 fn test_fixed_size_binary_encoding_applicable() {
592 assert!(!is_fixed_size_encoding_applicable(vec![vec![]]));
593
594 assert!(is_fixed_size_encoding_applicable(vec![vec![
595 Some("a"),
596 Some("b")
597 ]]));
598
599 assert!(!is_fixed_size_encoding_applicable(vec![vec![
600 Some("abc"),
601 Some("de")
602 ]]));
603
604 assert!(is_fixed_size_encoding_applicable(vec![vec![
605 Some("pqr"),
606 None
607 ]]));
608
609 assert!(!is_fixed_size_encoding_applicable(vec![vec![
610 Some("pqr"),
611 Some("")
612 ]]));
613
614 assert!(!is_fixed_size_encoding_applicable(vec![vec![
615 Some(""),
616 Some("")
617 ]]));
618 }
619
620 #[test]
621 fn test_fixed_size_binary_encoding_applicable_multiple_arrays() {
622 assert!(is_fixed_size_encoding_applicable(vec![
623 vec![Some("a"), Some("b")],
624 vec![Some("c"), Some("d")]
625 ]));
626
627 assert!(!is_fixed_size_encoding_applicable(vec![
628 vec![Some("ab"), Some("bc")],
629 vec![Some("c"), Some("d")]
630 ]));
631
632 assert!(!is_fixed_size_encoding_applicable(vec![
633 vec![Some("ab"), None],
634 vec![None, Some("d")]
635 ]));
636
637 assert!(is_fixed_size_encoding_applicable(vec![
638 vec![Some("a"), None],
639 vec![None, Some("d")]
640 ]));
641
642 assert!(!is_fixed_size_encoding_applicable(vec![
643 vec![Some(""), None],
644 vec![None, Some("")]
645 ]));
646
647 assert!(!is_fixed_size_encoding_applicable(vec![
648 vec![None, None],
649 vec![None, None]
650 ]));
651 }
652
653 fn verify_array_encoder(
654 array: ArrayRef,
655 field_meta: Option<HashMap<String, String>>,
656 expected_encoder: &str,
657 ) {
658 let encoding_strategy = ArrayStrategy;
659 let mut field = Field::new("test_field", array.data_type().clone(), true);
660 if let Some(field_meta) = field_meta {
661 field.set_metadata(field_meta);
662 }
663 let lance_field = lance_core::datatypes::Field::try_from(field).unwrap();
664 let encoder_result = encoding_strategy.create_array_encoder(&[array], &lance_field);
665 assert!(encoder_result.is_ok());
666 let encoder = encoder_result.unwrap();
667 assert_eq!(format!("{:?}", encoder).as_str(), expected_encoder);
668 }
669
670 #[test]
671 fn test_choose_encoder_for_zstd_compressed_string_field() {
672 verify_array_encoder(
673 Arc::new(StringArray::from(vec!["a", "bb", "ccc"])),
674 Some(HashMap::from([(
675 COMPRESSION_META_KEY.to_string(),
676 "zstd".to_string(),
677 )])),
678 "BinaryEncoder { indices_encoder: BasicEncoder { values_encoder: ValueEncoder }, compression_config: Some(CompressionConfig { scheme: Zstd, level: None }), buffer_compressor: Some(ZstdBufferCompressor { compression_level: 0 }) }",
679 );
680 }
681
682 #[test]
683 fn test_choose_encoder_for_zstd_compression_level() {
684 verify_array_encoder(
685 Arc::new(StringArray::from(vec!["a", "bb", "ccc"])),
686 Some(HashMap::from([
687 (COMPRESSION_META_KEY.to_string(), "zstd".to_string()),
688 (COMPRESSION_LEVEL_META_KEY.to_string(), "22".to_string()),
689 ])),
690 "BinaryEncoder { indices_encoder: BasicEncoder { values_encoder: ValueEncoder }, compression_config: Some(CompressionConfig { scheme: Zstd, level: Some(22) }), buffer_compressor: Some(ZstdBufferCompressor { compression_level: 22 }) }",
691 );
692 }
693}