diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py index 8e91d22c49fc..8f1e2d39ed00 100644 --- a/ibis/backends/tests/test_array.py +++ b/ibis/backends/tests/test_array.py @@ -1059,3 +1059,14 @@ def test_repr_timestamp_array(con, monkeypatch): assert ibis.options.default_backend is con expr = ibis.array(pd.date_range("2010-01-01", "2010-01-03", freq="D").tolist()) assert repr(expr) + + +@pytest.mark.notyet( + ["dask", "datafusion", "flink", "pandas", "polars"], + raises=com.OperationNotDefinedError, +) +def test_unnest_range(con): + expr = ibis.range(2).unnest().name("x").as_table().mutate({"y": 1.0}) + result = con.execute(expr) + expected = pd.DataFrame({"x": np.array([0, 1], dtype="int8"), "y": [1.0, 1.0]}) + tm.assert_frame_equal(result, expected) diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py index 6bfefccb1922..ad3b12ef363b 100644 --- a/ibis/backends/tests/test_generic.py +++ b/ibis/backends/tests/test_generic.py @@ -1741,3 +1741,16 @@ def test_simple_memtable_construct(con): # because memtables have a unique name per table per process, so smoke test # it assert str(ibis.to_sql(expr, dialect=con.name)).startswith("SELECT") + + +def test_select_mutate_with_dict(backend): + t = backend.functional_alltypes + expr = t.mutate({"a": 1.0}).select("a").limit(1) + + result = expr.execute() + expected = pd.DataFrame({"a": [1.0]}) + + backend.assert_frame_equal(result, expected) + + expr = t.select({"a": ibis.literal(1.0)}).limit(1) + backend.assert_frame_equal(result, expected) diff --git a/ibis/expr/types/relations.py b/ibis/expr/types/relations.py index 030442361f0a..a66dea273878 100644 --- a/ibis/expr/types/relations.py +++ b/ibis/expr/types/relations.py @@ -1800,17 +1800,24 @@ def mutate( import ibis.expr.analysis as an exprs = [] if exprs is None else util.promote_list(exprs) - exprs = itertools.chain( - itertools.chain.from_iterable( - util.promote_list(_ensure_expr(self, expr)) for expr in exprs - ), - ( - e.name(name) - for name, expr in mutations.items() - for e in util.promote_list(_ensure_expr(self, expr)) - ), + + new_exprs = [] + + for expr in exprs: + if isinstance(expr, Mapping): + new_exprs.extend( + _ensure_expr(self, val).name(name) for name, val in expr.items() + ) + else: + new_exprs.extend(util.promote_list(_ensure_expr(self, expr))) + + new_exprs.extend( + e.name(name) + for name, expr in mutations.items() + for e in util.promote_list(_ensure_expr(self, expr)) ) - mutation_exprs = an.get_mutation_exprs(list(exprs), self) + + mutation_exprs = an.get_mutation_exprs(new_exprs, self) return self.select(mutation_exprs) def select( @@ -1993,31 +2000,34 @@ def select( import ibis.expr.analysis as an from ibis.selectors import Selector - exprs = [ - e - for expr in exprs - for e in ( - expr.expand(self) - if isinstance(expr, Selector) - else map(self._ensure_expr, util.promote_list(expr)) - ) - ] - exprs.extend( + new_exprs = [] + + for expr in exprs: + if isinstance(expr, Selector): + new_exprs.extend(expr.expand(self)) + elif isinstance(expr, Mapping): + new_exprs.extend( + self._ensure_expr(value).name(name) for name, value in expr.items() + ) + else: + new_exprs.extend(map(self._ensure_expr, util.promote_list(expr))) + + new_exprs.extend( self._ensure_expr(expr).name(name) for name, expr in named_exprs.items() ) - if not exprs: + if not new_exprs: raise com.IbisTypeError( "You must select at least one column for a valid projection" ) - for ex in exprs: + for ex in new_exprs: if not isinstance(ex, Expr): raise com.IbisTypeError( "All arguments to `.select` must be coerceable to " f"expressions - got {type(ex)!r}" ) - op = an.Projector(self, exprs).get_result() + op = an.Projector(self, new_exprs).get_result() return op.to_expr() diff --git a/ibis/util.py b/ibis/util.py index d9e0bc1544e6..9b95a12d2c47 100644 --- a/ibis/util.py +++ b/ibis/util.py @@ -107,6 +107,8 @@ def promote_list(val: V | Sequence[V]) -> list[V]: """ if isinstance(val, list): return val + elif isinstance(val, dict): + return [val] elif is_iterable(val): return list(val) elif val is None: