Add value_type field to DataType (#93)

kylebarron · Jul 31, 2024 · 9d5f2d2 · 9d5f2d2
1 parent 2cdec15
commit 9d5f2d2
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/arro3-core/python/arro3/core/_core.pyi b/arro3-core/python/arro3/core/_core.pyi
@@ -181,6 +181,8 @@ class DataType:
     @property
     def num_fields(self) -> int:
         """The number of child fields."""
+    @property
+    def value_type(self) -> DataType | None: ...
     #################
     #### Constructors
     #################

diff --git a/pyo3-arrow/src/datatypes.rs b/pyo3-arrow/src/datatypes.rs
@@ -214,6 +214,23 @@ impl PyDataType {
         }
     }
 
+    #[getter]
+    fn value_type(&self, py: Python) -> PyResult<Option<PyObject>> {
+        match &self.0 {
+            DataType::List(value_field)
+            | DataType::LargeList(value_field)
+            | DataType::ListView(value_field)
+            | DataType::LargeListView(value_field)
+            | DataType::RunEndEncoded(_, value_field) => Ok(Some(
+                PyDataType::new(value_field.data_type().clone()).to_arro3(py)?,
+            )),
+            DataType::Dictionary(_key_type, value_type) => {
+                Ok(Some(PyDataType::new(*value_type.clone()).to_arro3(py)?))
+            }
+            _ => Ok(None),
+        }
+    }
+
     ///////////////////// Constructors
 
     #[classmethod]