Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions arro3-core/python/arro3/core/_data_type.pyi
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from typing import Literal, Sequence

from ._field import Field
from .types import ArrowSchemaExportable

class DataType:
Expand Down Expand Up @@ -73,6 +74,15 @@ class DataType:
@property
def value_type(self) -> DataType | None:
"""The child type, if it exists."""

@property
def value_field(self) -> Field | None:
"""The child field, if it exists. Only applicable to list types."""

@property
def fields(self) -> Sequence[Field]:
"""The inner fields, if they exists. Only applicable to Struct type."""

#################
#### Constructors
#################
Expand Down
29 changes: 28 additions & 1 deletion pyo3-arrow/src/datatypes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use pyo3::prelude::*;
use pyo3::types::{PyCapsule, PyTuple, PyType};

use crate::error::PyArrowResult;
use crate::export::Arro3DataType;
use crate::export::{Arro3DataType, Arro3Field};
use crate::ffi::from_python::utils::import_schema_pycapsule;
use crate::ffi::to_python::nanoarrow::to_nanoarrow_schema;
use crate::ffi::to_schema_pycapsule;
Expand Down Expand Up @@ -264,6 +264,33 @@ impl PyDataType {
}
}

#[getter]
fn value_field(&self) -> Option<Arro3Field> {
match &self.0 {
DataType::FixedSizeList(value_field, _)
| DataType::List(value_field)
| DataType::LargeList(value_field)
| DataType::ListView(value_field)
| DataType::LargeListView(value_field) => {
Some(PyField::new(value_field.clone()).into())
}
_ => None,
}
}

#[getter]
fn fields(&self) -> Option<Vec<Arro3Field>> {
match &self.0 {
DataType::Struct(fields) => Some(
fields
.into_iter()
.map(|f| PyField::new(f.clone()).into())
.collect::<Vec<_>>(),
),
_ => None,
}
}

///////////////////// Constructors

#[classmethod]
Expand Down
17 changes: 17 additions & 0 deletions tests/core/test_data_type.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,23 @@ def test_value_type_fixed_size_list_type():
assert list_dt.value_type == value_type


def test_value_field_list_type():
value_type = DataType.int8()
value_field = Field("inner", value_type, nullable=True)
list_dt = DataType.list(
value_field,
2,
)
assert list_dt.value_field == value_field


def test_fields_struct_type():
field_foo = Field("foo", DataType.int8(), nullable=True)
field_bar = Field("bar", DataType.string(), nullable=False)
struct_type = DataType.struct([field_foo, field_bar])
assert struct_type.fields == [field_foo, field_bar]


@pytest.mark.xfail
def test_list_data_type_construction_with_dt():
_ = DataType.list(DataType.int16())
Loading