Function test_dictionary_list_hydration

arrow-flight/src/encode.rs:970–1025

()

    #[tokio::test]
    async fn test_dictionary_list_hydration() {
        let mut builder = ListBuilder::new(StringDictionaryBuilder::<UInt16Type>::new());

        builder.append_value(vec![Some("a"), None, Some("b")]);

        let arr1 = builder.finish();

        builder.append_value(vec![Some("c"), None, Some("d")]);

        let arr2 = builder.finish();

        let schema = Arc::new(Schema::new(vec![Field::new_list(
            "dict_list",
            Field::new_dictionary("item", DataType::UInt16, DataType::Utf8, true),
            true,
        )]));

        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_list(
            "dict_list",
            Field::new_list_field(DataType::Utf8, true),
            true,
        )]);

        let expected_schema = Arc::new(expected_schema);

        let mut expected_arrays = vec![
            StringArray::from_iter(vec![Some("a"), None, Some("b")]),
            StringArray::from_iter(vec![Some("c"), None, Some("d")]),
        ]
        .into_iter();

        while let Some(decoded) = decoder.next().await {
            let decoded = decoded.unwrap();
            match decoded.payload {
                DecodedPayload::None => {}
                DecodedPayload::Schema(s) => assert_eq!(s, expected_schema),
                DecodedPayload::RecordBatch(b) => {
                    assert_eq!(b.schema(), expected_schema);
                    let expected_array = expected_arrays.next().unwrap();
                    let list_array =
                        downcast_array::<ListArray>(b.column_by_name("dict_list").unwrap());
                    let elem_array = downcast_array::<StringArray>(list_array.value(0).as_ref());

                    assert_eq!(elem_array, expected_array);
                }
            }
        }
    }

    #[tokio::test]