def test_log(diamonds: ir.TableExpr, base: int) -> None:
result_expr = diamonds >> log(nullif(X.carat, 0), base)
expected_expr = diamonds.carat.nullif(0).log(base)
assert result_expr.equals(expected_expr)
result_df = result_expr >> do()
expected_df = expected_expr.execute()
tm.assert_series_equal(result_df, expected_df)
def test_compound_expression(diamonds: ir.TableExpr) -> None:
expected = diamonds[diamonds.price * diamonds.price / 2.0 >= 100]
expected = expected.groupby('cut').aggregate([
expected.carat.max().name('max_carat'),
expected.carat.mean().name('mean_carat'),
expected.carat.min().name('min_carat'),
expected.x.count().name('n'),
expected.carat.std().name('std_carat'),
expected.carat.sum().name('sum_carat'),
expected.carat.var().name('var_carat'),
])
expected = expected.mutate(foo=expected.mean_carat,
bar=expected.var_carat).sort_by(
[ibis.desc('foo'), 'bar']).head()
result = (diamonds >> sift(X.price * X.price / 2.0 >= 100) >> groupby(
X.cut) >> summarize(
max_carat=max(X.carat),
mean_carat=mean(X.carat),
min_carat=min(X.carat),
n=n(X.x),
std_carat=std(X.carat),
sum_carat=sum(X.carat),
var_carat=var(X.carat),
) >> mutate(foo=X.mean_carat, bar=X.var_carat) >> sort_by(
desc(X.foo), X.bar) >> head(5))
assert result.equals(expected)
tm.assert_frame_equal(expected.execute(), result >> do())
def test_unary_string(diamonds: ir.TableExpr, func: Type[Unary]) -> None:
result = diamonds >> func(X.cut)
expected = getattr(diamonds.cut, func.__name__)()
assert result.equals(expected)
tm.assert_series_equal(result >> do(), expected.execute())
def test_round(diamonds: ir.TableExpr, places: int) -> None:
result = diamonds >> round(X.carat, places)
expected = diamonds.carat.round(places)
assert result.equals(expected)
tm.assert_series_equal(result >> do(), expected.execute())
def test_binary_math(diamonds: ir.TableExpr, func: Callable) -> None:
result = diamonds >> func(X.carat, X.z)
expected = func(diamonds.carat, diamonds.z)
assert result.equals(expected)
tm.assert_series_equal(result >> do(), expected.execute())
def test_unary_math(diamonds: ir.TableExpr, func: Type[Unary]) -> None:
result = diamonds >> func(cast(X.carat, to=dt.Decimal(19, 7)))
expected = getattr(diamonds.carat.cast(dt.Decimal(19, 7)), func.__name__)()
assert result.equals(expected)
tm.assert_series_equal(result >> do(), expected.execute())
def test_nunique(diamonds: ir.TableExpr, column: str) -> None:
result = diamonds >> nunique(X[column])
expected = diamonds[column].nunique()
assert result.equals(expected)
assert expected.execute() == result >> do()
def test_distinct(diamonds: ir.TableExpr, column: str) -> None:
result = diamonds >> distinct(X[column])
expected = diamonds[column].distinct()
assert result.equals(expected)
tm.assert_series_equal(expected.execute(), result >> do())
def test_cast(diamonds: ir.TableExpr, to: Union[str, dt.DataType]) -> None:
result = diamonds >> cast(X.carat + 1, to=to)
expected = (diamonds.carat + 1).cast(to)
assert result.equals(expected)
tm.assert_series_equal(expected.execute(), result >> do())
def test_transmute(diamonds: ir.TableExpr) -> None:
result = diamonds >> transmute(new_column=X.carat * 2)
expected = diamonds[[(diamonds.carat * 2).name('new_column')]]
assert result.equals(expected)
tm.assert_frame_equal(expected.execute(), result >> do())
def test_mutate(diamonds: ir.TableExpr) -> None:
result = diamonds >> mutate(new_column=X.carat + 1)
expected = diamonds.mutate(new_column=lambda x: x.carat + 1)
assert result.equals(expected)
tm.assert_frame_equal(expected.execute(), result >> do())
def test_simple_arithmetic(diamonds: ir.TableExpr) -> None:
result = diamonds >> mean(X.carat) + 1
expected = diamonds.carat.mean() + 1
assert result.equals(expected)
assert float(expected.execute()) == float(result >> do())
def test_do(diamonds: ir.TableExpr) -> None:
tm.assert_frame_equal(diamonds.execute(), diamonds >> do())
def test_pull(diamonds: ir.TableExpr, column: Union[str, int]) -> None:
result = diamonds >> X[column]
expected = diamonds[column]
assert result.equals(expected)
tm.assert_series_equal(expected.execute(), result >> do())