Skip to content

Commit

Permalink
feat(api): make struct metadata more convenient to access
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Jul 5, 2022
1 parent fab4393 commit 3fd9bd8
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 1 deletion.
19 changes: 18 additions & 1 deletion ibis/expr/types/structs.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@

import collections
import itertools
from typing import TYPE_CHECKING, Iterable, Mapping
from typing import TYPE_CHECKING, Iterable, Mapping, Sequence

from cached_property import cached_property
from public import public

from ibis import util
from ibis.expr.types.generic import Column, Scalar, Value, literal
from ibis.expr.types.typing import V

Expand Down Expand Up @@ -78,6 +80,21 @@ def __getitem__(self, name: str) -> ir.Value:

return ops.StructField(self, name).to_expr().name(name)

@cached_property
def names(self) -> Sequence[str]:
"""Return the field names of the struct."""
return self.type().names

@cached_property
def types(self) -> Sequence[dt.DataType]:
"""Return the field types of the struct."""
return self.type().types

@cached_property
def fields(self) -> Mapping[str, dt.DataType]:
"""Return a mapping from field name to field type of the struct."""
return util.frozendict(self.type().pairs)

def destructure(self) -> DestructValue:
"""Destructure `self` into a `DestructValue`.
Expand Down
20 changes: 20 additions & 0 deletions ibis/tests/expr/test_value_exprs.py
Original file line number Diff line number Diff line change
Expand Up @@ -1524,3 +1524,23 @@ def test_non_null_with_null_precedence(expr_fn, expected_type):
t = ibis.table(dict(a="int64", b="!string"), name="t")
expr = expr_fn(t)
assert expr.type() == expected_type


@pytest.mark.parametrize(
("name", "expected"),
[
("names", ("a", "b", "c")),
("types", (dt.int8, dt.string, dt.Array(dt.Array(dt.float64)))),
(
"fields",
ibis.util.frozendict(
a=dt.int8,
b=dt.string,
c=dt.Array(dt.Array(dt.float64)),
),
),
],
)
def test_struct_names_types_fields(name, expected):
s = ibis.struct(dict(a=1, b="2", c=[[1.0], [], [None]]))
assert getattr(s, name) == expected

0 comments on commit 3fd9bd8

Please sign in to comment.