Remove intermediate variable

0x26res · Sep 24, 2024 · 05f3e4b · 05f3e4b
1 parent 71e665b
commit 05f3e4b
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 13 deletions.
diff --git a/python/test/test_ptars.py b/python/test/test_ptars.py
@@ -116,10 +116,16 @@ def test_arrow_to_proto(pool):
         [
             pa.array([33, 42, 17], pa.int32()),
             pa.array([133, 142, 117], pa.int64()),
+            pa.array([True, False, None], pa.bool_()),
+            pa.array(["ABC", "", None], pa.utf8()),
+            pa.array([b"ABC", b"", None], pa.binary()),
         ],
         [
             "int32_value",
             "int64_value",
+            "bool_value",
+            "string_value",
+            "bytes_value",
         ],
     )
     handler = pool.get_for_message(ExampleMessage.DESCRIPTOR)
@@ -128,15 +134,15 @@ def test_arrow_to_proto(pool):
     assert array.type == pa.binary()
     assert len(array) == 3
 
-    payload = array.to_pylist()[0]
-    assert len(payload) == 5
-    assert ExampleMessage.FromString(payload) == ExampleMessage(
-        int32_value=33, int64_value=133
-    )
-
     messages = [ExampleMessage.FromString(b.as_py()) for b in array]
     assert messages == [
-        ExampleMessage(int32_value=33, int64_value=133),
+        ExampleMessage(
+            int32_value=33,
+            int64_value=133,
+            bool_value=True,
+            string_value="ABC",
+            bytes_value=b"ABC",
+        ),
         ExampleMessage(int32_value=42, int64_value=142),
         ExampleMessage(int32_value=17, int64_value=117),
     ]

diff --git a/src/lib.rs b/src/lib.rs
@@ -556,8 +556,10 @@ fn set_primitive<P: ArrowPrimitiveType>(
     field_descriptor: &FieldDescriptor,
     rvb_creator: &dyn Fn(P::Native) -> ReflectValueBox,
 ) {
-    let specific_array: &PrimitiveArray<P> = array.as_any().downcast_ref().unwrap();
-    specific_array
+    array
+        .as_any()
+        .downcast_ref::<PrimitiveArray<P>>()
+        .unwrap()
         .iter()
         .enumerate()
         .for_each(|(index, value)| match value {
@@ -616,9 +618,52 @@ fn extract_singular_array(
             field_descriptor,
             &ReflectValueBox::F64,
         ),
-        RuntimeType::Bool => {} // BooleanType doesn't implement primitive type
-        RuntimeType::String => {}
-        RuntimeType::VecU8 => {}
+        RuntimeType::Bool => {
+            // BooleanType doesn't implement primitive type
+            array
+                .as_any()
+                .downcast_ref::<BooleanArray>()
+                .unwrap()
+                .iter()
+                .enumerate()
+                .for_each(|(index, value)| match value {
+                    None => {}
+                    Some(x) => {
+                        let element: &mut dyn MessageDyn =
+                            messages.get_mut(index).unwrap().as_mut();
+                        field_descriptor
+                            .set_singular_field(&mut *element, ReflectValueBox::Bool(x));
+                    }
+                })
+        }
+        RuntimeType::String => array
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap()
+            .iter()
+            .enumerate()
+            .for_each(|(index, value)| match value {
+                None => {}
+                Some(x) => {
+                    let element: &mut dyn MessageDyn = messages.get_mut(index).unwrap().as_mut();
+                    field_descriptor
+                        .set_singular_field(&mut *element, ReflectValueBox::String(x.to_string()));
+                }
+            }),
+        RuntimeType::VecU8 => array
+            .as_any()
+            .downcast_ref::<BinaryArray>()
+            .unwrap()
+            .iter()
+            .enumerate()
+            .for_each(|(index, value)| match value {
+                None => {}
+                Some(x) => {
+                    let element: &mut dyn MessageDyn = messages.get_mut(index).unwrap().as_mut();
+                    field_descriptor
+                        .set_singular_field(&mut *element, ReflectValueBox::Bytes(x.to_vec()));
+                }
+            }),
         RuntimeType::Enum(_) => {}
         RuntimeType::Message(_) => {}
     }
@@ -629,7 +674,6 @@ fn extract_array(
     field_descriptor: &FieldDescriptor,
     messages: &mut Vec<Box<dyn MessageDyn>>,
 ) {
-    println!("!!!  {}", field_descriptor.name());
     match field_descriptor.runtime_field_type() {
         RuntimeFieldType::Singular(x) => {
             extract_singular_array(&array, &field_descriptor, messages, &x)