def update_spec(expr: ibis.Expr, spec: dict): """ Takes in an ibis expression and a spec, updating the spec and returning a new ibis expr """ original_expr = expr # iterate through transforms and move as many as we can into the ibis expression # logic modified from # https://github.com/vega/vega-lite-transforms2sql/blob/3b360144305a6cec79792036049e8a920e4d2c9e/transforms2sql.ts#L7 for transform in spec.get("transform", []): groupby = transform.pop("groupby", None) if groupby: all_fields_exist = all( [field in expr.columns for field in groupby]) if not all_fields_exist: transform["groupby"] = groupby # we referenced a field that isnt in the expression because it was an aggregate we coudnt process continue expr = expr.groupby(groupby) aggregate = transform.pop("aggregate", None) if aggregate: expr = expr.aggregate([ vl_aggregate_to_grouping_expr(original_expr, a) for a in aggregate ]) filter_ = transform.pop("filter", None) if filter_: # https://vega.github.io/vega-lite/docs/filter.html#field-predicate field = filter_["field"] field_expr = original_expr[field] if "range" in filter_: min, max = filter_["range"] preds = [field_expr >= min, field_expr <= max] elif "equal" in filter_: preds = [field_expr == filter_["equal"]] elif "gt" in filter_: preds = [field_expr > filter_["gt"]] elif "lt" in filter_: preds = [field_expr < filter_["lt"]] elif "lte" in filter_: preds = [field_expr <= filter_["lte"]] elif "gte" in filter_: preds = [field_expr >= filter_["gte"]] else: # put filter back if we cant transform itt transform["filter"] = filter_ continue expr = expr.filter(preds) # remove empty transforms spec["transform"] = [i for i in spec.get("transform", []) if i] # remove key if empty if not spec["transform"]: del spec["transform"] return expr
def filter(transform: dict, expr: ibis.Expr) -> ibis.Expr: calc = transform["expr"] return expr.filter(eval_vegajs(calc, expr))