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

lance_encoding/encodings/logical/
map.rs
