Skip to content

Commit

Permalink
feat: make map(), struct(), array() idempotent
Browse files Browse the repository at this point in the history
Before `array(array([1,2]))` would error.
Now it doesn't

This is in prep for the next commit.
But it also is just a nice behavior, where
you can now call these functions as many times as you want and always get back something that makes sense.

Also, fix a typing bug: map() can accept ArrayValues, not just ArrayColumns
  • Loading branch information
NickCrews committed Mar 15, 2024
1 parent 220085e commit b6d983a
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 8 deletions.
6 changes: 4 additions & 2 deletions ibis/expr/types/arrays.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,13 @@
from public import public

import ibis.expr.operations as ops
import ibis.expr.types as ir
from ibis.common.deferred import Deferred, deferrable
from ibis.expr.types.generic import Column, Scalar, Value

if TYPE_CHECKING:
from collections.abc import Iterable

import ibis.expr.types as ir
from ibis.expr.types.typing import V

import ibis.common.exceptions as com
Expand Down Expand Up @@ -1081,7 +1081,7 @@ def __getitem__(self, index: int | ir.IntegerValue | slice) -> ir.Column:

@public
@deferrable
def array(values: Iterable[V]) -> ArrayValue:
def array(values: ArrayValue | Iterable[V]) -> ArrayValue:
"""Create an array expression.
If any values are [column expressions](../concepts/datatypes.qmd) the
Expand Down Expand Up @@ -1131,4 +1131,6 @@ def array(values: Iterable[V]) -> ArrayValue:
│ [3, 47] │
└──────────────────────┘
"""
if isinstance(values, ir.ArrayValue):
return values
return ops.Array(tuple(values)).to_expr()
10 changes: 5 additions & 5 deletions ibis/expr/types/maps.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,13 @@
from public import public

import ibis.expr.operations as ops
import ibis.expr.types as ir
from ibis.common.deferred import deferrable
from ibis.expr.types.generic import Column, Scalar, Value

if TYPE_CHECKING:
from collections.abc import Iterable, Mapping

import ibis.expr.types as ir
from ibis.expr.types.arrays import ArrayColumn


@public
class MapValue(Value):
Expand Down Expand Up @@ -435,8 +433,8 @@ def __getitem__(self, key: ir.Value) -> ir.Column:
@public
@deferrable
def map(
keys: Iterable[Any] | Mapping[Any, Any] | ArrayColumn,
values: Iterable[Any] | ArrayColumn | None = None,
keys: Iterable[Any] | Mapping[Any, Any] | ir.ArrayValue | MapValue,
values: Iterable[Any] | ir.ArrayValue | None = None,
) -> MapValue:
"""Create a MapValue.
Expand Down Expand Up @@ -486,6 +484,8 @@ def map(
│ {'b': 3} │
└──────────────────────┘
"""
if isinstance(keys, MapValue):
return keys
if values is None:
keys, values = tuple(keys.keys()), tuple(keys.values())
return ops.Map(keys, values).to_expr()
7 changes: 6 additions & 1 deletion ibis/expr/types/structs.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
@public
@deferrable
def struct(
value: Iterable[tuple[str, V]] | Mapping[str, V],
value: Iterable[tuple[str, V]] | Mapping[str, V] | StructValue,
type: str | dt.DataType | None = None,
) -> StructValue:
"""Create a struct expression.
Expand Down Expand Up @@ -75,6 +75,11 @@ def struct(
"""
import ibis.expr.operations as ops

if isinstance(value, StructValue):
if type is not None:
return value.cast(type)
return value

fields = dict(value)
if any(isinstance(value, Value) for value in fields.values()):
names = tuple(fields.keys())
Expand Down

0 comments on commit b6d983a

Please sign in to comment.