lance_encoding/encodings/logical/
map.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{ops::Range, sync::Arc};
5
6use arrow_array::{Array, ArrayRef, ListArray, MapArray};
7use arrow_schema::DataType;
8use futures::future::BoxFuture;
9use lance_arrow::deepcopy::deep_copy_nulls;
10use lance_arrow::list::ListArrayExt;
11use lance_core::{Error, Result};
12
13use crate::{
14    decoder::{
15        DecodedArray, FilterExpression, ScheduledScanLine, SchedulerContext,
16        StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
17        StructuralSchedulingJob,
18    },
19    encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers},
20    repdef::RepDefBuilder,
21};
22
23/// A structural encoder for map fields
24///
25/// Map in Arrow is represented as List<Struct<key, value>>
26/// The map's offsets are added to the rep/def builder
27/// and the map's entries (struct array) are passed to the child encoder
28pub struct MapStructuralEncoder {
29    keep_original_array: bool,
30    child: Box<dyn FieldEncoder>,
31}
32
33impl MapStructuralEncoder {
34    pub fn new(keep_original_array: bool, child: Box<dyn FieldEncoder>) -> Self {
35        Self {
36            keep_original_array,
37            child,
38        }
39    }
40}
41
42impl FieldEncoder for MapStructuralEncoder {
43    fn maybe_encode(
44        &mut self,
45        array: ArrayRef,
46        external_buffers: &mut OutOfLineBuffers,
47        mut repdef: RepDefBuilder,
48        row_number: u64,
49        num_rows: u64,
50    ) -> Result<Vec<EncodeTask>> {
51        let map_array = array
52            .as_any()
53            .downcast_ref::<MapArray>()
54            .expect("MapEncoder used for non-map data");
55
56        // Add offsets to RepDefBuilder to handle nullability and list structure
57        let has_garbage_values = if self.keep_original_array {
58            repdef.add_offsets(map_array.offsets().clone(), array.nulls().cloned())
59        } else {
60            repdef.add_offsets(map_array.offsets().clone(), deep_copy_nulls(array.nulls()))
61        };
62
63        // MapArray is physically a ListArray, so convert and use ListArrayExt
64        let list_array: ListArray = map_array.clone().into();
65        let entries = if has_garbage_values {
66            list_array.filter_garbage_nulls().trimmed_values()
67        } else {
68            list_array.trimmed_values()
69        };
70
71        self.child
72            .maybe_encode(entries, external_buffers, repdef, row_number, num_rows)
73    }
74
75    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
76        self.child.flush(external_buffers)
77    }
78
79    fn num_columns(&self) -> u32 {
80        self.child.num_columns()
81    }
82
83    fn finish(
84        &mut self,
85        external_buffers: &mut OutOfLineBuffers,
86    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
87        self.child.finish(external_buffers)
88    }
89}
90
91#[derive(Debug)]
92pub struct StructuralMapScheduler {
93    child: Box<dyn StructuralFieldScheduler>,
94}
95
96impl StructuralMapScheduler {
97    pub fn new(child: Box<dyn StructuralFieldScheduler>) -> Self {
98        Self { child }
99    }
100}
101
102impl StructuralFieldScheduler for StructuralMapScheduler {
103    fn schedule_ranges<'a>(
104        &'a self,
105        ranges: &[Range<u64>],
106        filter: &FilterExpression,
107    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
108        let child = self.child.schedule_ranges(ranges, filter)?;
109
110        Ok(Box::new(StructuralMapSchedulingJob::new(child)))
111    }
112
113    fn initialize<'a>(
114        &'a mut self,
115        filter: &'a FilterExpression,
116        context: &'a SchedulerContext,
117    ) -> BoxFuture<'a, Result<()>> {
118        self.child.initialize(filter, context)
119    }
120}
121
122/// Scheduling job for map data
123///
124/// Scheduling is handled by the child encoder (struct) and nothing special
125/// happens here, similar to list.
126#[derive(Debug)]
127struct StructuralMapSchedulingJob<'a> {
128    child: Box<dyn StructuralSchedulingJob + 'a>,
129}
130
131impl<'a> StructuralMapSchedulingJob<'a> {
132    fn new(child: Box<dyn StructuralSchedulingJob + 'a>) -> Self {
133        Self { child }
134    }
135}
136
137impl StructuralSchedulingJob for StructuralMapSchedulingJob<'_> {
138    fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
139        self.child.schedule_next(context)
140    }
141}
142
143#[derive(Debug)]
144pub struct StructuralMapDecoder {
145    child: Box<dyn StructuralFieldDecoder>,
146    data_type: DataType,
147}
148
149impl StructuralMapDecoder {
150    pub fn new(child: Box<dyn StructuralFieldDecoder>, data_type: DataType) -> Self {
151        Self { child, data_type }
152    }
153}
154
155impl StructuralFieldDecoder for StructuralMapDecoder {
156    fn accept_page(&mut self, child: crate::decoder::LoadedPageShard) -> Result<()> {
157        self.child.accept_page(child)
158    }
159
160    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
161        let child_task = self.child.drain(num_rows)?;
162        Ok(Box::new(StructuralMapDecodeTask::new(
163            child_task,
164            self.data_type.clone(),
165        )))
166    }
167
168    fn data_type(&self) -> &DataType {
169        &self.data_type
170    }
171}
172
173#[derive(Debug)]
174struct StructuralMapDecodeTask {
175    child_task: Box<dyn StructuralDecodeArrayTask>,
176    data_type: DataType,
177}
178
179impl StructuralMapDecodeTask {
180    fn new(child_task: Box<dyn StructuralDecodeArrayTask>, data_type: DataType) -> Self {
181        Self {
182            child_task,
183            data_type,
184        }
185    }
186}
187
188impl StructuralDecodeArrayTask for StructuralMapDecodeTask {
189    fn decode(self: Box<Self>) -> Result<DecodedArray> {
190        let DecodedArray {
191            array,
192            mut repdef,
193            data_size,
194        } = self.child_task.decode()?;
195
196        // Decode the offsets from RepDef
197        let (offsets, validity) = repdef.unravel_offsets::<i32>()?;
198
199        // Extract the entries field and keys_sorted from the map data type
200        let (entries_field, keys_sorted) = match &self.data_type {
201            DataType::Map(field, keys_sorted) => {
202                if *keys_sorted {
203                    return Err(Error::not_supported_source(
204                        "Map type decoder does not support keys_sorted=true now"
205                            .to_string()
206                            .into(),
207                    ));
208                }
209                (field.clone(), *keys_sorted)
210            }
211            _ => {
212                return Err(Error::schema(
213                    "Map decoder did not have a map field".to_string(),
214                ));
215            }
216        };
217
218        // Convert the decoded array to StructArray
219        let entries = array
220            .as_any()
221            .downcast_ref::<arrow_array::StructArray>()
222            .ok_or_else(|| Error::schema("Map entries should be a StructArray".to_string()))?
223            .clone();
224
225        // Build the MapArray from offsets, entries, validity, and keys_sorted
226        let map_array = MapArray::try_new(entries_field, offsets, entries, validity, keys_sorted)
227            .map_err(|error| Error::invalid_input_source(error.to_string().into()))?;
228
229        Ok(DecodedArray {
230            array: Arc::new(map_array),
231            repdef,
232            data_size,
233        })
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use std::{collections::HashMap, sync::Arc};
240
241    use arrow_array::{
242        Array, Int32Array, MapArray, StringArray, StructArray,
243        builder::{Int32Builder, MapBuilder, StringBuilder},
244    };
245    use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
246    use arrow_schema::{DataType, Field, Fields};
247
248    use crate::decoder::{DecodedArray, StructuralDecodeArrayTask};
249    use crate::encoder::{ColumnIndexSequence, EncodingOptions};
250    use crate::encodings::logical::primitive::sparse::{
251        SparseCountSet, SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan,
252        SparseValidityMeaning, SparseValiditySet,
253    };
254    use crate::repdef::{CompositeRepDefUnraveler, RepDefUnraveler};
255    use crate::testing::{
256        TestCases, TestEncoding, check_round_trip_encoding_of_data, test_encoding_strategy,
257    };
258    use arrow_schema::Field as ArrowField;
259    use lance_core::datatypes::Field as LanceField;
260
261    use super::StructuralMapDecodeTask;
262
263    fn make_map_type(key_type: DataType, value_type: DataType) -> DataType {
264        // Note: Arrow MapBuilder uses "keys" and "values" as field names (plural)
265        let entries = Field::new(
266            "entries",
267            DataType::Struct(Fields::from(vec![
268                Field::new("keys", key_type, false),
269                Field::new("values", value_type, true),
270            ])),
271            false,
272        );
273        DataType::Map(Arc::new(entries), false)
274    }
275
276    #[derive(Debug)]
277    struct StaticMapEntriesTask {
278        entries: StructArray,
279        repdef: CompositeRepDefUnraveler,
280    }
281
282    impl StructuralDecodeArrayTask for StaticMapEntriesTask {
283        fn decode(self: Box<Self>) -> lance_core::Result<DecodedArray> {
284            let Self { entries, repdef } = *self;
285            Ok(DecodedArray {
286                array: Arc::new(entries),
287                repdef,
288                data_size: 0,
289            })
290        }
291    }
292
293    #[test]
294    fn malformed_sparse_map_entries_return_invalid_input() {
295        let entry_fields = Fields::from(vec![
296            Field::new("keys", DataType::Int32, false),
297            Field::new("values", DataType::Int32, true),
298        ]);
299        let entries = StructArray::try_new(
300            entry_fields.clone(),
301            vec![
302                Arc::new(Int32Array::from(vec![1])),
303                Arc::new(Int32Array::from(vec![2])),
304            ],
305            Some(NullBuffer::from(vec![false])),
306        )
307        .unwrap();
308        let validity = SparseValiditySet {
309            meaning: SparseValidityMeaning::NullPositions,
310            positions: SparsePositionSet::Empty,
311        };
312        let plan = SparseStructuralPlan {
313            layers: vec![SparseStructuralLayerPlan::List {
314                num_slots: 1,
315                num_child_slots: 1,
316                non_empty_positions: SparsePositionSet::All { len: 1 },
317                counts: SparseCountSet::Constant { value: 1, len: 1 },
318                validity,
319            }],
320            num_items: 1,
321            num_visible_items: 1,
322        };
323        let child_task = StaticMapEntriesTask {
324            entries,
325            repdef: CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new_sparse(plan)]),
326        };
327        let map_type = DataType::Map(
328            Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)),
329            false,
330        );
331
332        let Err(err) =
333            Box::new(StructuralMapDecodeTask::new(Box::new(child_task), map_type)).decode()
334        else {
335            panic!("expected malformed map entries to be rejected");
336        };
337        assert!(matches!(err, lance_core::Error::InvalidInput { .. }));
338    }
339
340    #[test_log::test(tokio::test)]
341    async fn test_simple_map() {
342        // Create a simple Map<String, Int32>
343        let string_builder = StringBuilder::new();
344        let int_builder = Int32Builder::new();
345        let mut map_builder = MapBuilder::new(None, string_builder, int_builder);
346
347        // Map 1: {"key1": 10, "key2": 20}
348        map_builder.keys().append_value("key1");
349        map_builder.values().append_value(10);
350        map_builder.keys().append_value("key2");
351        map_builder.values().append_value(20);
352        map_builder.append(true).unwrap();
353
354        // Map 2: {"key3": 30}
355        map_builder.keys().append_value("key3");
356        map_builder.values().append_value(30);
357        map_builder.append(true).unwrap();
358
359        let map_array = map_builder.finish();
360
361        let test_cases = TestCases::default()
362            .with_range(0..2)
363            .with_u32_structural_encodings();
364
365        check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new())
366            .await;
367    }
368
369    #[test_log::test(tokio::test)]
370    async fn test_empty_maps() {
371        // Test maps with empty entries
372        let string_builder = StringBuilder::new();
373        let int_builder = Int32Builder::new();
374        let mut map_builder = MapBuilder::new(None, string_builder, int_builder);
375
376        // Map 1: {"a": 1}
377        map_builder.keys().append_value("a");
378        map_builder.values().append_value(1);
379        map_builder.append(true).unwrap();
380
381        // Map 2: {} (empty)
382        map_builder.append(true).unwrap();
383
384        // Map 3: null
385        map_builder.append(false).unwrap();
386
387        // Map 4: {} (empty)
388        map_builder.append(true).unwrap();
389
390        let map_array = map_builder.finish();
391
392        let test_cases = TestCases::default()
393            .with_range(0..4)
394            .with_indices(vec![1])
395            .with_indices(vec![2])
396            .with_u32_structural_encodings();
397
398        check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new())
399            .await;
400    }
401
402    #[test_log::test(tokio::test)]
403    async fn test_map_with_null_values() {
404        // Test Map<String, Int32> with null values
405        let string_builder = StringBuilder::new();
406        let int_builder = Int32Builder::new();
407        let mut map_builder = MapBuilder::new(None, string_builder, int_builder);
408
409        // Map 1: {"key1": 10, "key2": null}
410        map_builder.keys().append_value("key1");
411        map_builder.values().append_value(10);
412        map_builder.keys().append_value("key2");
413        map_builder.values().append_null();
414        map_builder.append(true).unwrap();
415
416        // Map 2: {"key3": null}
417        map_builder.keys().append_value("key3");
418        map_builder.values().append_null();
419        map_builder.append(true).unwrap();
420
421        let map_array = map_builder.finish();
422
423        let test_cases = TestCases::default()
424            .with_range(0..2)
425            .with_indices(vec![0])
426            .with_indices(vec![1])
427            .with_u32_structural_encodings();
428
429        check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new())
430            .await;
431    }
432
433    #[test_log::test(tokio::test)]
434    async fn test_map_in_struct() {
435        // Test Struct containing Map
436        // Struct<id: Int32, properties: Map<String, String>>
437
438        let string_key_builder = StringBuilder::new();
439        let string_val_builder = StringBuilder::new();
440        let mut map_builder = MapBuilder::new(None, string_key_builder, string_val_builder);
441
442        // First struct: id=1, properties={"name": "Alice", "city": "NYC"}
443        map_builder.keys().append_value("name");
444        map_builder.values().append_value("Alice");
445        map_builder.keys().append_value("city");
446        map_builder.values().append_value("NYC");
447        map_builder.append(true).unwrap();
448
449        // Second struct: id=2, properties={"name": "Bob"}
450        map_builder.keys().append_value("name");
451        map_builder.values().append_value("Bob");
452        map_builder.append(true).unwrap();
453
454        // Third struct: id=3, properties=null
455        map_builder.append(false).unwrap();
456
457        let map_array = Arc::new(map_builder.finish());
458        let id_array = Arc::new(Int32Array::from(vec![1, 2, 3]));
459
460        let struct_array = StructArray::new(
461            Fields::from(vec![
462                Field::new("id", DataType::Int32, false),
463                Field::new(
464                    "properties",
465                    make_map_type(DataType::Utf8, DataType::Utf8),
466                    true,
467                ),
468            ]),
469            vec![id_array, map_array],
470            None,
471        );
472
473        let test_cases = TestCases::default()
474            .with_range(0..3)
475            .with_indices(vec![0, 2])
476            .with_u32_structural_encodings();
477
478        check_round_trip_encoding_of_data(
479            vec![Arc::new(struct_array)],
480            &test_cases,
481            HashMap::new(),
482        )
483        .await;
484    }
485
486    #[test_log::test(tokio::test)]
487    async fn test_map_in_nullable_struct() {
488        // Test Struct<Map> where null struct rows have garbage map entries.
489        // The encoder must filter these garbage entries before encoding.
490        let entries_fields = Fields::from(vec![
491            Field::new("keys", DataType::Utf8, false),
492            Field::new("values", DataType::Int32, true),
493        ]);
494        let entries_field = Arc::new(Field::new(
495            "entries",
496            DataType::Struct(entries_fields.clone()),
497            false,
498        ));
499        let map_entries = StructArray::new(
500            entries_fields,
501            vec![
502                Arc::new(StringArray::from(vec!["a", "garbage", "b"])),
503                Arc::new(Int32Array::from(vec![1, 999, 2])),
504            ],
505            None,
506        );
507        // map0: {"a": 1}, map1 (garbage): {"garbage": 999}, map2: {"b": 2}
508        let map_array: Arc<dyn Array> = Arc::new(MapArray::new(
509            entries_field,
510            OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2, 3])),
511            map_entries,
512            None, // No nulls at map level - nulls come from struct
513            false,
514        ));
515
516        let struct_array = StructArray::new(
517            Fields::from(vec![
518                Field::new("id", DataType::Int32, true),
519                Field::new("props", map_array.data_type().clone(), true),
520            ]),
521            vec![
522                Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)])),
523                map_array,
524            ],
525            Some(NullBuffer::from(vec![true, false, true])), // Middle row is null
526        );
527
528        let test_cases = TestCases::default()
529            .with_range(0..3)
530            .with_u32_structural_encodings();
531
532        check_round_trip_encoding_of_data(
533            vec![Arc::new(struct_array)],
534            &test_cases,
535            HashMap::new(),
536        )
537        .await;
538    }
539
540    #[test_log::test(tokio::test)]
541    async fn test_list_of_maps() {
542        // Test List<Map<String, Int32>>
543        use arrow_array::builder::ListBuilder;
544
545        let string_builder = StringBuilder::new();
546        let int_builder = Int32Builder::new();
547        let map_builder = MapBuilder::new(None, string_builder, int_builder);
548        let mut list_builder = ListBuilder::new(map_builder);
549
550        // List 1: [{"a": 1}, {"b": 2}]
551        list_builder.values().keys().append_value("a");
552        list_builder.values().values().append_value(1);
553        list_builder.values().append(true).unwrap();
554
555        list_builder.values().keys().append_value("b");
556        list_builder.values().values().append_value(2);
557        list_builder.values().append(true).unwrap();
558
559        list_builder.append(true);
560
561        // List 2: [{"c": 3}]
562        list_builder.values().keys().append_value("c");
563        list_builder.values().values().append_value(3);
564        list_builder.values().append(true).unwrap();
565
566        list_builder.append(true);
567
568        // List 3: [] (empty list)
569        list_builder.append(true);
570
571        let list_array = list_builder.finish();
572
573        let test_cases = TestCases::default()
574            .with_range(0..3)
575            .with_indices(vec![0, 2])
576            .with_u32_structural_encodings();
577
578        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
579            .await;
580    }
581
582    #[test_log::test(tokio::test)]
583    async fn test_nested_map() {
584        // Test Map<String, Map<String, Int32>>
585        // This is more complex as we need to build nested maps manually
586
587        // Build inner maps first
588        let inner_string_builder = StringBuilder::new();
589        let inner_int_builder = Int32Builder::new();
590        let mut inner_map_builder1 = MapBuilder::new(None, inner_string_builder, inner_int_builder);
591
592        // Inner map 1: {"x": 10}
593        inner_map_builder1.keys().append_value("x");
594        inner_map_builder1.values().append_value(10);
595        inner_map_builder1.append(true).unwrap();
596
597        // Inner map 2: {"y": 20, "z": 30}
598        inner_map_builder1.keys().append_value("y");
599        inner_map_builder1.values().append_value(20);
600        inner_map_builder1.keys().append_value("z");
601        inner_map_builder1.values().append_value(30);
602        inner_map_builder1.append(true).unwrap();
603
604        let inner_maps = Arc::new(inner_map_builder1.finish());
605
606        // Build outer map keys
607        let outer_keys = Arc::new(StringArray::from(vec!["key1", "key2"]));
608
609        // Build outer map structure
610        let entries_struct = StructArray::new(
611            Fields::from(vec![
612                Field::new("key", DataType::Utf8, false),
613                Field::new(
614                    "value",
615                    make_map_type(DataType::Utf8, DataType::Int32),
616                    true,
617                ),
618            ]),
619            vec![outer_keys, inner_maps],
620            None,
621        );
622
623        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2]));
624        let entries_field = Field::new("entries", entries_struct.data_type().clone(), false);
625
626        let outer_map = MapArray::new(
627            Arc::new(entries_field),
628            offsets,
629            entries_struct,
630            None,
631            false,
632        );
633
634        let test_cases = TestCases::default()
635            .with_range(0..1)
636            .with_u32_structural_encodings();
637
638        check_round_trip_encoding_of_data(vec![Arc::new(outer_map)], &test_cases, HashMap::new())
639            .await;
640    }
641
642    #[test_log::test(tokio::test)]
643    async fn test_map_different_key_types() {
644        // Test Map<Int32, String> (integer keys)
645        let int_builder = Int32Builder::new();
646        let string_builder = StringBuilder::new();
647        let mut map_builder = MapBuilder::new(None, int_builder, string_builder);
648
649        // Map 1: {1: "one", 2: "two"}
650        map_builder.keys().append_value(1);
651        map_builder.values().append_value("one");
652        map_builder.keys().append_value(2);
653        map_builder.values().append_value("two");
654        map_builder.append(true).unwrap();
655
656        // Map 2: {3: "three"}
657        map_builder.keys().append_value(3);
658        map_builder.values().append_value("three");
659        map_builder.append(true).unwrap();
660
661        let map_array = map_builder.finish();
662
663        let test_cases = TestCases::default()
664            .with_range(0..2)
665            .with_indices(vec![0, 1])
666            .with_u32_structural_encodings();
667
668        check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new())
669            .await;
670    }
671
672    #[test_log::test(tokio::test)]
673    async fn test_map_with_extreme_sizes() {
674        // Test maps with large number of entries
675        let string_builder = StringBuilder::new();
676        let int_builder = Int32Builder::new();
677        let mut map_builder = MapBuilder::new(None, string_builder, int_builder);
678
679        // Create a map with many entries
680        for i in 0..100 {
681            map_builder.keys().append_value(format!("key{}", i));
682            map_builder.values().append_value(i);
683        }
684        map_builder.append(true).unwrap();
685
686        // Create a second map with no entries
687        map_builder.append(true).unwrap();
688
689        let map_array = map_builder.finish();
690
691        let test_cases = TestCases::default()
692            .with_range(0..2)
693            .with_u32_structural_encodings();
694
695        check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new())
696            .await;
697    }
698
699    #[test_log::test(tokio::test)]
700    async fn test_map_all_null() {
701        // Test map where all entries are null
702        let string_builder = StringBuilder::new();
703        let int_builder = Int32Builder::new();
704        let mut map_builder = MapBuilder::new(None, string_builder, int_builder);
705
706        // All null maps
707        map_builder.append(false).unwrap(); // null
708        map_builder.append(false).unwrap(); // null
709
710        let map_array = map_builder.finish();
711
712        let test_cases = TestCases::default()
713            .with_range(0..2)
714            .with_u32_structural_encodings();
715
716        check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new())
717            .await;
718    }
719
720    #[test_log::test(tokio::test)]
721    async fn test_map_encoder_keep_original_array_scenarios() {
722        // Test scenarios that highlight the difference between keep_original_array=true/false
723        // This test focuses on round-trip behavior which should be equivalent in both cases
724        let string_builder = StringBuilder::new();
725        let int_builder = Int32Builder::new();
726        let mut map_builder = MapBuilder::new(None, string_builder, int_builder);
727
728        // Create a map with mixed null and non-null values to test both scenarios
729        // Map 1: {"key1": 10, "key2": null}
730        map_builder.keys().append_value("key1");
731        map_builder.values().append_value(10);
732        map_builder.keys().append_value("key2");
733        map_builder.values().append_null();
734        map_builder.append(true).unwrap();
735
736        // Map 2: null
737        map_builder.append(false).unwrap();
738
739        // Map 3: {"key3": 30}
740        map_builder.keys().append_value("key3");
741        map_builder.values().append_value(30);
742        map_builder.append(true).unwrap();
743
744        let map_array = map_builder.finish();
745
746        let test_cases = TestCases::default()
747            .with_range(0..3)
748            .with_indices(vec![0, 1, 2])
749            .with_u32_structural_encodings();
750
751        // This test ensures that regardless of the internal keep_original_array setting,
752        // the end-to-end behavior produces equivalent results
753        check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new())
754            .await;
755    }
756
757    #[test]
758    fn test_map_not_supported_write_in_v2_1() {
759        // Create a map field using Arrow Field first, then convert to Lance Field
760        let map_arrow_field = ArrowField::new(
761            "map_field",
762            make_map_type(DataType::Utf8, DataType::Int32),
763            true,
764        );
765        let map_field = LanceField::try_from(&map_arrow_field).unwrap();
766
767        // Test encoder: Try to create encoder with V2_1 version - should fail
768        let encoder_strategy = test_encoding_strategy(TestEncoding::StructuralU16);
769        let mut column_index = ColumnIndexSequence::default();
770        let options = EncodingOptions::default();
771
772        let encoder_result = crate::testing::create_test_field_encoder(
773            encoder_strategy.as_ref(),
774            &map_field,
775            &mut column_index,
776            &options,
777        );
778
779        assert!(
780            encoder_result.is_err(),
781            "Map type should not be supported in V2_1 for encoder"
782        );
783        let Err(encoder_err) = encoder_result else {
784            panic!("Expected error but got Ok")
785        };
786
787        let encoder_err_msg = format!("{}", encoder_err);
788        assert!(
789            encoder_err_msg.contains("not enabled by the selected file format"),
790            "unexpected encoder error: {encoder_err_msg}"
791        );
792        assert!(
793            encoder_err_msg.contains("Map data type"),
794            "Encoder error message should mention Map data type, got: {}",
795            encoder_err_msg
796        );
797    }
798}
lance_encoding/encodings/logical/map.rs

lance_encoding/encodings/logical/
map.rs
