Skip to content

Commit

Permalink
fix(struct-column): make ops.StructColumn dshape depend on its input
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Aug 7, 2023
1 parent dbe7d4e commit 7086d58
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 3 deletions.
2 changes: 1 addition & 1 deletion ibis/backends/clickhouse/compiler/values.py
Expand Up @@ -409,7 +409,7 @@ def _log(op, **kw):
@translate_val.register(tuple)
def _node_list(op, **kw):
values = ", ".join(map(_sql, map(partial(translate_val, **kw), op)))
return f"({values})"
return f"tuple({values})"


def _interval_format(op):
Expand Down
22 changes: 22 additions & 0 deletions ibis/backends/tests/test_struct.py
Expand Up @@ -88,3 +88,25 @@ def test_struct_column(alltypes, df):
name="s",
)
tm.assert_series_equal(result, expected)


@pytest.mark.notimpl(["dask", "pandas", "postgres", "polars"])
def test_collect_into_struct(alltypes):
from ibis import _

t = alltypes
expr = (
t[_.string_col.isin(("0", "1"))]
.group_by(group="string_col")
.agg(
val=lambda t: ibis.struct(
dict(key=t.bigint_col.collect().cast("!array<int64>"))
)
)
)
result = expr.execute()
assert result.shape == (2, 2)
assert set(result.group) == {"0", "1"}
val = result.val
assert len(val.loc[result.group == "0"].iat[0]["key"]) == 730
assert len(val.loc[result.group == "1"].iat[0]["key"]) == 730
3 changes: 1 addition & 2 deletions ibis/expr/operations/structs.py
Expand Up @@ -2,7 +2,6 @@

from public import public

import ibis.expr.datashape as ds
import ibis.expr.datatypes as dt
import ibis.expr.rules as rlz
from ibis.common.annotations import attribute
Expand Down Expand Up @@ -33,7 +32,7 @@ class StructColumn(Value):
names: VarTuple[str]
values: VarTuple[Value]

shape = ds.columnar
shape = rlz.shape_like("values")

@attribute.default
def dtype(self) -> dt.DataType:
Expand Down
19 changes: 19 additions & 0 deletions ibis/expr/operations/tests/test_structs.py
@@ -0,0 +1,19 @@
from __future__ import annotations

import ibis
import ibis.expr.datashape as ds
import ibis.expr.datatypes as dt
import ibis.expr.operations as ops


def test_struct_column_shape():
one = ops.Literal(1, dtype=dt.int64)
op = ops.StructColumn(names=("a",), values=(one,))

assert op.shape == ds.scalar

col = ops.TableColumn(
ops.UnboundTable(schema=ibis.schema(dict(a="int64")), name="t"), "a"
)
op = ops.StructColumn(names=("a",), values=(col,))
assert op.shape == ds.columnar

0 comments on commit 7086d58

Please sign in to comment.