Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
@select.register(LazyTbl)
def _(__data, *args, **kwargs):
# see https://stackoverflow.com/questions/25914329/rearrange-columns-in-sqlalchemy-select-object
last_op = __data.last_op
if isinstance(last_op, sql.Select):
columns = {c.key: c for c in last_op.inner_columns}
# same as for DataFrame
colnames = Series(list(columns))
vl = VarList()
evaluated = (strip_symbolic(arg)(vl) if callable(arg) else arg for arg in args)
od = var_select(colnames, *evaluated)
col_list = []
for k,v in od.items():
col = columns[k]
col_list.append(col if v is None else col.label(v))
@mutate.register(LazyTbl)
def _(__data, **kwargs):
# Cases
# - work with group by
# - window functions
# TODO: can't re-use select, for example if it's following a select() that renames
# track labeled columns in set
#
last_op = __data.last_op
labs = set(k for k,v in last_op.columns.items() if isinstance(v, sql.elements.Label))
sel = last_op
# TODO: could use copy of last_op plus append_column
for colname, func in kwargs.items():
inner_cols = lift_inner_cols(sel)
replace_col = colname in sel.columns
@count.register(LazyTbl)
def _(__data, *args, sort = False):
# TODO: need group_by to do this
last_op = __data.last_op
__data.append_op(last_op)
@arrange.register(LazyTbl)
def _(__data, *args):
last_op = __data.last_op
cols = lift_inner_cols(last_op)
sort_cols = []
for arg in args:
# simple named column
if isinstance(arg, str):
sort_cols.append(cols[arg])
# an expression
elif callable(arg):
f, asc = _call_strip_ascending(arg)
col_op = f(cols) if asc else f(cols).desc()
sort_cols.append(col_op)
else:
raise NotImplementedError("Must be string or callable")
@summarize.register(LazyTbl)
def _(__data, **kwargs):
# https://stackoverflow.com/questions/14754994/why-is-sqlalchemy-count-much-slower-than-the-raw-query
pass
@filter.register(LazyTbl)
def _(__data, *args, **kwargs):
# TODO: use inner_columns to prevent nesting selects?
last_op = __data.last_op
cols = lift_inner_cols(last_op)
conds = [arg(cols) if callable(arg) else arg for arg in args]
bool_clause = sql.and_(*conds)
if isinstance(last_op, sql.Select):
return __data.append_op(last_op.where(bool_clause))
return __data.append_op(sql.Select(['*'], whereclause = bool_clause))
mtcars data.
Source: Henderson and Velleman (1981), Building multiple regression models interactively. Biometrics, 37, 391–411.
--- Original DataFrame docs below ---
""" + mtcars.__doc__
# cars ------------------------------------------------------------------------
cars = mtcars[["cyl", "mpg", "hp"]]
# cars_sql --------------------------------------------------------------------
import siuba.sql.utils as _sql_utils
from siuba.sql import LazyTbl as _LazyTbl
cars_sql = _LazyTbl(
_sql_utils.mock_sqlalchemy_engine("postgresql"),
"cars",
["cyl", "mpg", "hp"]
)