Function test_dictionary_map_hydration

arrow-flight/src/encode.rs:1451–1554

()

    #[tokio::test]
    async fn test_dictionary_map_hydration() {
        let mut builder = MapBuilder::new(
            None,
            StringDictionaryBuilder::<UInt16Type>::new(),
            StringDictionaryBuilder::<UInt16Type>::new(),
        );

        // {"k1":"a","k2":null,"k3":"b"}
        builder.keys().append_value("k1");
        builder.values().append_value("a");
        builder.keys().append_value("k2");
        builder.values().append_null();
        builder.keys().append_value("k3");
        builder.values().append_value("b");
        builder.append(true).unwrap();

        let arr1 = builder.finish();

        // {"k1":"c","k2":null,"k3":"d"}
        builder.keys().append_value("k1");
        builder.values().append_value("c");
        builder.keys().append_value("k2");
        builder.values().append_null();
        builder.keys().append_value("k3");
        builder.values().append_value("d");
        builder.append(true).unwrap();

        let arr2 = builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_map(
            "dict_map",
            "entries",
            Field::new_dictionary("keys", DataType::UInt16, DataType::Utf8, false),
            Field::new_dictionary("values", DataType::UInt16, DataType::Utf8, true),
            false,
            false,
        )]));

        let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr1)]).unwrap();
        let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr2)]).unwrap();

        let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]);

        let encoder = FlightDataEncoderBuilder::default().build(stream);

        let mut decoder = FlightDataDecoder::new(encoder);
        let expected_schema = Schema::new(vec![Field::new_map(
            "dict_map",
            "entries",
            Field::new("keys", DataType::Utf8, false),
            Field::new("values", DataType::Utf8, true),
            false,
            false,
        )]);

        let expected_schema = Arc::new(expected_schema);

        // Builder without dictionary fields