From 81b00b51148b1c02c58751d96cccafa67e540220 Mon Sep 17 00:00:00 2001 From: fjetter Date: Wed, 6 Dec 2023 15:53:54 +0100 Subject: [PATCH] Add tuple expr --- dask_expr/_expr.py | 11 +++++++++++ dask_expr/tests/test_collection.py | 14 +++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py index 2d1e5a914..1a4b38e29 100644 --- a/dask_expr/_expr.py +++ b/dask_expr/_expr.py @@ -2394,6 +2394,17 @@ class Pos(Unaryop): _operator_repr = "+" +class Tuple(Expr): + def __getitem__(self, other): + return self.operands[other] + + def __len__(self): + return len(self.operands) + + def __iter__(self): + return iter(self.operands) + + class Partitions(Expr): """Select one or more partitions""" diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py index 88dcca71a..b8611e500 100644 --- a/dask_expr/tests/test_collection.py +++ b/dask_expr/tests/test_collection.py @@ -19,7 +19,7 @@ to_numeric, to_timedelta, ) -from dask_expr._expr import are_co_aligned +from dask_expr._expr import Tuple, are_co_aligned from dask_expr._reductions import Len from dask_expr._shuffle import Shuffle from dask_expr.datasets import timeseries @@ -1565,3 +1565,15 @@ def test_items(df, pdf): for (expect_name, expect_col), (actual_name, actual_col) in zip(expect, actual): assert expect_name == actual_name assert_eq(expect_col, actual_col) + + +def test_combine_expr_with_tuple(pdf): + ddf1 = from_pandas(pdf, npartitions=2) + 1 + ddf2 = from_pandas(pdf, npartitions=3) + 2 + + t = Tuple(ddf1.expr, ddf2.expr) + assert t[0]._name == ddf1._name + assert t[0].optimize()._name == t.optimize()[0]._name + + assert t[1]._name == ddf2._name + assert t[1].optimize()._name == t.optimize()[1]._name