feat(dataflow): model production SQL flow patterns
This commit is contained in:
@@ -498,6 +498,16 @@ def _join_rows(
|
||||
for left_row in left_rows:
|
||||
key = _join_key(left_row, left_keys)
|
||||
matches = right_index.get(key, ()) if key is not None else ()
|
||||
if join_type == "semi":
|
||||
if matches:
|
||||
output.append(dict(left_row))
|
||||
_guard_intermediate_size(output, node_id=node_id)
|
||||
continue
|
||||
if join_type == "anti":
|
||||
if not matches:
|
||||
output.append(dict(left_row))
|
||||
_guard_intermediate_size(output, node_id=node_id)
|
||||
continue
|
||||
if matches:
|
||||
for right_index_value, right_row in matches:
|
||||
matched_right.add(right_index_value)
|
||||
@@ -688,6 +698,41 @@ def _expression_rows(
|
||||
return output
|
||||
|
||||
|
||||
def _calculation_rows(
|
||||
rows: list[dict[str, Any]],
|
||||
config: dict[str, Any],
|
||||
*,
|
||||
node_id: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
calculations: list[tuple[str, Any]] = []
|
||||
try:
|
||||
calculations = [
|
||||
(
|
||||
str(item["target_column"]),
|
||||
parse_expression(str(item["expression"])),
|
||||
)
|
||||
for item in config["calculations"]
|
||||
]
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise PipelineExecutionError(
|
||||
f"Cannot prepare calculated columns: {exc}",
|
||||
node_id=node_id,
|
||||
) from exc
|
||||
output: list[dict[str, Any]] = []
|
||||
for row in rows:
|
||||
result = dict(row)
|
||||
for target, parsed in calculations:
|
||||
try:
|
||||
result[target] = evaluate_expression(parsed, result)
|
||||
except (ArithmeticError, TypeError, ValueError) as exc:
|
||||
raise PipelineExecutionError(
|
||||
f"Cannot calculate {target!r}: {exc}",
|
||||
node_id=node_id,
|
||||
) from exc
|
||||
output.append(result)
|
||||
return output
|
||||
|
||||
|
||||
def _convert_rows(
|
||||
rows: list[dict[str, Any]],
|
||||
config: dict[str, Any],
|
||||
@@ -1165,6 +1210,65 @@ def _sort_rows(rows: list[dict[str, Any]], config: dict[str, Any]) -> list[dict[
|
||||
return result
|
||||
|
||||
|
||||
def _rank_rows(
|
||||
rows: list[dict[str, Any]],
|
||||
config: dict[str, Any],
|
||||
) -> list[dict[str, Any]]:
|
||||
partition_by = [str(item) for item in config.get("partition_by", [])]
|
||||
order_by = list(config["order_by"])
|
||||
partitions: dict[tuple[Any, ...], list[tuple[int, dict[str, Any]]]] = (
|
||||
defaultdict(list)
|
||||
)
|
||||
for index, row in enumerate(rows):
|
||||
key = tuple(_hashable(row.get(column)) for column in partition_by)
|
||||
partitions[key].append((index, row))
|
||||
|
||||
ranks: dict[int, int] = {}
|
||||
method = str(config.get("method", "row_number"))
|
||||
for partition in partitions.values():
|
||||
ordered = list(partition)
|
||||
for field_config in reversed(order_by):
|
||||
column = str(field_config["column"])
|
||||
reverse = field_config.get("direction", "asc") == "desc"
|
||||
concrete = [
|
||||
item for item in ordered if item[1].get(column) is not None
|
||||
]
|
||||
nulls = [
|
||||
item for item in ordered if item[1].get(column) is None
|
||||
]
|
||||
concrete.sort(
|
||||
key=lambda item: _sortable_value(item[1][column]),
|
||||
reverse=reverse,
|
||||
)
|
||||
ordered = [*concrete, *nulls]
|
||||
|
||||
previous_values: tuple[Any, ...] | None = None
|
||||
current_rank = 0
|
||||
dense_rank = 0
|
||||
for position, (source_index, row) in enumerate(ordered, start=1):
|
||||
values = tuple(
|
||||
_hashable(row.get(str(field["column"])))
|
||||
for field in order_by
|
||||
)
|
||||
if previous_values is None or values != previous_values:
|
||||
current_rank = position
|
||||
dense_rank += 1
|
||||
previous_values = values
|
||||
ranks[source_index] = (
|
||||
position
|
||||
if method == "row_number"
|
||||
else dense_rank
|
||||
if method == "dense_rank"
|
||||
else current_rank
|
||||
)
|
||||
|
||||
target = str(config["target_column"])
|
||||
return [
|
||||
{**row, target: ranks[index]}
|
||||
for index, row in enumerate(rows)
|
||||
]
|
||||
|
||||
|
||||
def _sortable_value(value: Any) -> tuple[str, Any]:
|
||||
if isinstance(value, (int, float, Decimal, str)):
|
||||
return type(value).__name__, value
|
||||
@@ -1336,6 +1440,13 @@ def _register_executors() -> None:
|
||||
node_id=context.node.id,
|
||||
)
|
||||
),
|
||||
"calculate": lambda context: OperatorExecutionResult(
|
||||
rows=_calculation_rows(
|
||||
context.input_rows,
|
||||
context.node.config,
|
||||
node_id=context.node.id,
|
||||
)
|
||||
),
|
||||
"convert": lambda context: OperatorExecutionResult(
|
||||
rows=_convert_rows(
|
||||
context.input_rows,
|
||||
@@ -1356,6 +1467,9 @@ def _register_executors() -> None:
|
||||
"sort": lambda context: OperatorExecutionResult(
|
||||
rows=_sort_rows(context.input_rows, context.node.config)
|
||||
),
|
||||
"window.rank": lambda context: OperatorExecutionResult(
|
||||
rows=_rank_rows(context.input_rows, context.node.config)
|
||||
),
|
||||
"limit": lambda context: OperatorExecutionResult(
|
||||
rows=context.input_rows[: int(context.node.config["count"])]
|
||||
),
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import calendar
|
||||
import operator
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from typing import Any, Callable, Literal
|
||||
|
||||
@@ -36,6 +39,12 @@ class ParsedExpression:
|
||||
return self.expression.sql(dialect="duckdb")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _IntervalValue:
|
||||
amount: int
|
||||
unit: str
|
||||
|
||||
|
||||
_LEAF_TYPES = (
|
||||
exp.Column,
|
||||
exp.Identifier,
|
||||
@@ -43,6 +52,8 @@ _LEAF_TYPES = (
|
||||
exp.Null,
|
||||
exp.Boolean,
|
||||
exp.DataType,
|
||||
exp.Interval,
|
||||
exp.Var,
|
||||
)
|
||||
_BINARY_TYPES = (
|
||||
exp.Add,
|
||||
@@ -50,6 +61,7 @@ _BINARY_TYPES = (
|
||||
exp.Mul,
|
||||
exp.Div,
|
||||
exp.Mod,
|
||||
exp.DPipe,
|
||||
exp.EQ,
|
||||
exp.NEQ,
|
||||
exp.GT,
|
||||
@@ -60,6 +72,14 @@ _BINARY_TYPES = (
|
||||
exp.Or,
|
||||
exp.Is,
|
||||
)
|
||||
_PREDICATE_TYPES = (
|
||||
exp.Between,
|
||||
exp.ILike,
|
||||
exp.In,
|
||||
exp.Like,
|
||||
exp.RegexpFullMatch,
|
||||
exp.RegexpLike,
|
||||
)
|
||||
_UNARY_TYPES = (exp.Not, exp.Neg, exp.Paren)
|
||||
_FUNCTION_TYPES = (
|
||||
exp.Lower,
|
||||
@@ -73,11 +93,30 @@ _FUNCTION_TYPES = (
|
||||
exp.Replace,
|
||||
exp.Substring,
|
||||
exp.Cast,
|
||||
exp.ConcatWs,
|
||||
exp.DateAdd,
|
||||
exp.DateDiff,
|
||||
exp.DateSub,
|
||||
exp.Day,
|
||||
exp.Extract,
|
||||
exp.Greatest,
|
||||
exp.Least,
|
||||
exp.Month,
|
||||
exp.Nullif,
|
||||
exp.Pad,
|
||||
exp.RegexpReplace,
|
||||
exp.SplitPart,
|
||||
exp.TimeToStr,
|
||||
exp.TimestampTrunc,
|
||||
exp.ToChar,
|
||||
exp.Translate,
|
||||
exp.Year,
|
||||
)
|
||||
_CONTROL_TYPES = (exp.Case, exp.If)
|
||||
_ALLOWED_TYPES = (
|
||||
*_LEAF_TYPES,
|
||||
*_BINARY_TYPES,
|
||||
*_PREDICATE_TYPES,
|
||||
*_UNARY_TYPES,
|
||||
*_FUNCTION_TYPES,
|
||||
*_CONTROL_TYPES,
|
||||
@@ -182,6 +221,81 @@ def _evaluate_is(expression: exp.Expression, row: dict[str, Any]) -> bool:
|
||||
return left is right if right is None else left == right
|
||||
|
||||
|
||||
def _evaluate_like(expression: exp.Expression, row: dict[str, Any]) -> bool:
|
||||
value = _evaluate(expression.this, row)
|
||||
pattern = _evaluate(expression.expression, row)
|
||||
if value is None or pattern is None:
|
||||
return False
|
||||
flags = re.IGNORECASE if isinstance(expression, exp.ILike) else 0
|
||||
return (
|
||||
re.fullmatch(
|
||||
_sql_like_pattern(str(pattern)),
|
||||
str(value),
|
||||
flags=flags,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
|
||||
def _sql_like_pattern(pattern: str) -> str:
|
||||
pieces: list[str] = []
|
||||
escaped = False
|
||||
for character in pattern:
|
||||
if escaped:
|
||||
pieces.append(re.escape(character))
|
||||
escaped = False
|
||||
elif character == "\\":
|
||||
escaped = True
|
||||
elif character == "%":
|
||||
pieces.append(".*")
|
||||
elif character == "_":
|
||||
pieces.append(".")
|
||||
else:
|
||||
pieces.append(re.escape(character))
|
||||
if escaped:
|
||||
pieces.append(re.escape("\\"))
|
||||
return "".join(pieces)
|
||||
|
||||
|
||||
def _evaluate_in(expression: exp.Expression, row: dict[str, Any]) -> bool:
|
||||
value = _hashable_expression_value(_evaluate(expression.this, row))
|
||||
return value in {
|
||||
_hashable_expression_value(_evaluate(item, row))
|
||||
for item in expression.expressions
|
||||
}
|
||||
|
||||
|
||||
def _hashable_expression_value(value: Any) -> Any:
|
||||
if isinstance(value, (dict, list, set, tuple)):
|
||||
return repr(value)
|
||||
return value
|
||||
|
||||
|
||||
def _evaluate_between(expression: exp.Expression, row: dict[str, Any]) -> bool:
|
||||
value = _evaluate(expression.this, row)
|
||||
low = _evaluate(expression.args["low"], row)
|
||||
high = _evaluate(expression.args["high"], row)
|
||||
if value is None or low is None or high is None:
|
||||
return False
|
||||
return low <= value <= high
|
||||
|
||||
|
||||
def _evaluate_regexp_match(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> bool:
|
||||
value = _evaluate(expression.this, row)
|
||||
pattern = _evaluate(expression.expression, row)
|
||||
if value is None or pattern is None:
|
||||
return False
|
||||
operation = (
|
||||
re.fullmatch
|
||||
if isinstance(expression, exp.RegexpFullMatch)
|
||||
else re.search
|
||||
)
|
||||
return operation(str(pattern), str(value)) is not None
|
||||
|
||||
|
||||
def _evaluate_binary(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
@@ -197,6 +311,16 @@ def _evaluate_binary(
|
||||
|
||||
|
||||
def _evaluate_arithmetic(expression: exp.Expression, row: dict[str, Any]) -> Any:
|
||||
if isinstance(expression, (exp.Add, exp.Sub)):
|
||||
left = _evaluate(expression.this, row)
|
||||
right = _evaluate(expression.expression, row)
|
||||
if isinstance(right, _IntervalValue):
|
||||
amount = -right.amount if isinstance(expression, exp.Sub) else right.amount
|
||||
return _add_interval(left, amount, right.unit)
|
||||
if isinstance(left, _IntervalValue):
|
||||
if isinstance(expression, exp.Sub):
|
||||
raise ValueError("An interval cannot be subtracted from an interval.")
|
||||
return _add_interval(right, left.amount, left.unit)
|
||||
return _evaluate_binary(
|
||||
expression,
|
||||
row,
|
||||
@@ -232,6 +356,27 @@ def _evaluate_string_function(
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_trim(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
if value is None:
|
||||
return None
|
||||
characters_expression = expression.args.get("expression")
|
||||
characters = (
|
||||
str(_evaluate(characters_expression, row))
|
||||
if characters_expression is not None
|
||||
else None
|
||||
)
|
||||
position = str(expression.args.get("position") or "BOTH").upper()
|
||||
if position == "LEADING":
|
||||
return str(value).lstrip(characters)
|
||||
if position == "TRAILING":
|
||||
return str(value).rstrip(characters)
|
||||
return str(value).strip(characters)
|
||||
|
||||
|
||||
def _evaluate_length(expression: exp.Expression, row: dict[str, Any]) -> int | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
return None if value is None else len(value)
|
||||
@@ -286,6 +431,426 @@ def _evaluate_substring(expression: exp.Expression, row: dict[str, Any]) -> Any:
|
||||
return str(value)[start : start + int(_evaluate(length, row))]
|
||||
|
||||
|
||||
def _evaluate_pad(expression: exp.Expression, row: dict[str, Any]) -> str | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
if value is None:
|
||||
return None
|
||||
target_length = int(_evaluate(expression.expression, row))
|
||||
if target_length < 0:
|
||||
raise ValueError("Padding length cannot be negative.")
|
||||
source = str(value)
|
||||
if len(source) >= target_length:
|
||||
return source[:target_length]
|
||||
fill_expression = expression.args.get("fill_pattern")
|
||||
fill = (
|
||||
str(_evaluate(fill_expression, row))
|
||||
if fill_expression is not None
|
||||
else " "
|
||||
)
|
||||
if not fill:
|
||||
raise ValueError("Padding fill text cannot be empty.")
|
||||
padding_length = target_length - len(source)
|
||||
padding = (fill * ((padding_length // len(fill)) + 1))[:padding_length]
|
||||
return (
|
||||
f"{padding}{source}"
|
||||
if expression.args.get("is_left")
|
||||
else f"{source}{padding}"
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_split_part(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
delimiter = _evaluate(expression.args["delimiter"], row)
|
||||
part_index = int(_evaluate(expression.args["part_index"], row))
|
||||
if value is None or delimiter is None:
|
||||
return None
|
||||
if not delimiter:
|
||||
raise ValueError("SPLIT_PART delimiter cannot be empty.")
|
||||
if part_index == 0:
|
||||
raise ValueError("SPLIT_PART index is one-based.")
|
||||
parts = str(value).split(str(delimiter))
|
||||
index = part_index - 1 if part_index > 0 else part_index
|
||||
return parts[index] if -len(parts) <= index < len(parts) else ""
|
||||
|
||||
|
||||
def _evaluate_translate(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
source = _evaluate(expression.args["from_"], row)
|
||||
target = _evaluate(expression.args["to"], row)
|
||||
if value is None or source is None or target is None:
|
||||
return None
|
||||
target_text = str(target)
|
||||
translation = {
|
||||
ord(character): (
|
||||
target_text[index]
|
||||
if index < len(target_text)
|
||||
else None
|
||||
)
|
||||
for index, character in enumerate(str(source))
|
||||
}
|
||||
return str(value).translate(translation)
|
||||
|
||||
|
||||
def _evaluate_regexp_replace(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
pattern = _evaluate(expression.expression, row)
|
||||
replacement = _evaluate(expression.args["replacement"], row)
|
||||
if value is None or pattern is None or replacement is None:
|
||||
return None
|
||||
modifiers_expression = expression.args.get("modifiers")
|
||||
modifiers = (
|
||||
str(_evaluate(modifiers_expression, row))
|
||||
if modifiers_expression is not None
|
||||
else ""
|
||||
)
|
||||
flags = re.IGNORECASE if "i" in modifiers else 0
|
||||
count = 0 if "g" in modifiers else 1
|
||||
return re.sub(
|
||||
str(pattern),
|
||||
str(replacement),
|
||||
str(value),
|
||||
count=count,
|
||||
flags=flags,
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_concat_ws(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str:
|
||||
arguments = list(expression.expressions)
|
||||
separator_value = _evaluate(arguments[0], row) if arguments else None
|
||||
separator = "" if separator_value is None else str(separator_value)
|
||||
return separator.join(
|
||||
str(value)
|
||||
for argument in arguments[1:]
|
||||
if (value := _evaluate(argument, row)) is not None
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_pipe(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str | None:
|
||||
left = _evaluate(expression.this, row)
|
||||
right = _evaluate(expression.expression, row)
|
||||
return None if left is None or right is None else f"{left}{right}"
|
||||
|
||||
|
||||
def _evaluate_nullif(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> Any:
|
||||
value = _evaluate_child(expression, row)
|
||||
comparison = _evaluate(expression.expression, row)
|
||||
return None if value == comparison else value
|
||||
|
||||
|
||||
def _evaluate_extreme(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> Any:
|
||||
values = [
|
||||
value
|
||||
for argument in (expression.this, *expression.expressions)
|
||||
if (value := _evaluate(argument, row)) is not None
|
||||
]
|
||||
if not values:
|
||||
return None
|
||||
return (
|
||||
max(values)
|
||||
if isinstance(expression, exp.Greatest)
|
||||
else min(values)
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_time_to_str(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
format_expression = expression.args.get("format")
|
||||
format_text = (
|
||||
_evaluate(format_expression, row)
|
||||
if format_expression is not None
|
||||
else None
|
||||
)
|
||||
if value is None or format_text is None:
|
||||
return None
|
||||
return _format_temporal(value, str(format_text), postgres=False)
|
||||
|
||||
|
||||
def _evaluate_to_char(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> str | None:
|
||||
value = _evaluate_child(expression, row)
|
||||
format_expression = expression.args.get("format")
|
||||
format_text = (
|
||||
_evaluate(format_expression, row)
|
||||
if format_expression is not None
|
||||
else None
|
||||
)
|
||||
if value is None or format_text is None:
|
||||
return None
|
||||
if not isinstance(value, (date, datetime)) and not _looks_temporal(value):
|
||||
raise ValueError(
|
||||
"TO_CHAR currently supports date and date-time values only."
|
||||
)
|
||||
return _format_temporal(value, str(format_text), postgres=True)
|
||||
|
||||
|
||||
def _format_temporal(
|
||||
value: Any,
|
||||
format_text: str,
|
||||
*,
|
||||
postgres: bool,
|
||||
) -> str:
|
||||
if not postgres:
|
||||
return _temporal_value(value).strftime(format_text)
|
||||
translated = format_text
|
||||
for source, target in (
|
||||
("TMMonth", "%B"),
|
||||
("TMMon", "%b"),
|
||||
("TMYYYY", "%Y"),
|
||||
("YYYY", "%Y"),
|
||||
("YYY", "%Y"),
|
||||
("YY", "%y"),
|
||||
("MM", "%m"),
|
||||
("DD", "%d"),
|
||||
("HH24", "%H"),
|
||||
("MI", "%M"),
|
||||
("SS", "%S"),
|
||||
):
|
||||
translated = translated.replace(source, target)
|
||||
return _temporal_value(value).strftime(translated)
|
||||
|
||||
|
||||
def _looks_temporal(value: Any) -> bool:
|
||||
text = str(value or "").strip()
|
||||
if not text:
|
||||
return False
|
||||
try:
|
||||
_temporal_value(text)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
|
||||
|
||||
def _evaluate_extract(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> int:
|
||||
value = _temporal_value(_evaluate(expression.expression, row))
|
||||
unit = _date_unit(expression.this)
|
||||
return _temporal_part(value, unit)
|
||||
|
||||
|
||||
def _evaluate_temporal_part(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> int:
|
||||
value = _temporal_value(_evaluate_child(expression, row))
|
||||
unit = {
|
||||
exp.Year: "year",
|
||||
exp.Month: "month",
|
||||
exp.Day: "day",
|
||||
}[type(expression)]
|
||||
return _temporal_part(value, unit)
|
||||
|
||||
|
||||
def _temporal_part(value: date | datetime, unit: str) -> int:
|
||||
mapping = {
|
||||
"year": value.year,
|
||||
"quarter": ((value.month - 1) // 3) + 1,
|
||||
"month": value.month,
|
||||
"week": int(value.strftime("%V")),
|
||||
"day": value.day,
|
||||
"dow": value.weekday(),
|
||||
"doy": int(value.strftime("%j")),
|
||||
"hour": value.hour if isinstance(value, datetime) else 0,
|
||||
"minute": value.minute if isinstance(value, datetime) else 0,
|
||||
"second": value.second if isinstance(value, datetime) else 0,
|
||||
}
|
||||
if unit not in mapping:
|
||||
raise ValueError(f"Unsupported date part {unit!r}.")
|
||||
return mapping[unit]
|
||||
|
||||
|
||||
def _evaluate_timestamp_trunc(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> date | datetime:
|
||||
value = _temporal_value(_evaluate_child(expression, row))
|
||||
unit = _date_unit(expression.args["unit"])
|
||||
if unit == "year":
|
||||
return value.replace(month=1, day=1, **_time_reset(value))
|
||||
if unit == "quarter":
|
||||
month = ((value.month - 1) // 3) * 3 + 1
|
||||
return value.replace(month=month, day=1, **_time_reset(value))
|
||||
if unit == "month":
|
||||
return value.replace(day=1, **_time_reset(value))
|
||||
if unit == "week":
|
||||
result = value - timedelta(days=value.weekday())
|
||||
return result.replace(**_time_reset(result))
|
||||
if unit == "day":
|
||||
return value.replace(**_time_reset(value))
|
||||
if not isinstance(value, datetime):
|
||||
raise ValueError(f"Cannot truncate a date to {unit}.")
|
||||
if unit == "hour":
|
||||
return value.replace(minute=0, second=0, microsecond=0)
|
||||
if unit == "minute":
|
||||
return value.replace(second=0, microsecond=0)
|
||||
if unit == "second":
|
||||
return value.replace(microsecond=0)
|
||||
raise ValueError(f"Unsupported date truncation unit {unit!r}.")
|
||||
|
||||
|
||||
def _time_reset(value: date | datetime) -> dict[str, int]:
|
||||
return (
|
||||
{"hour": 0, "minute": 0, "second": 0, "microsecond": 0}
|
||||
if isinstance(value, datetime)
|
||||
else {}
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_date_add(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> date | datetime:
|
||||
value = _evaluate_child(expression, row)
|
||||
amount, unit = _interval_value(expression.expression, row)
|
||||
if isinstance(expression, (exp.DateSub, exp.Sub)):
|
||||
amount = -amount
|
||||
return _add_interval(value, amount, unit)
|
||||
|
||||
|
||||
def _add_interval(
|
||||
value: Any,
|
||||
amount: int,
|
||||
unit: str,
|
||||
) -> date | datetime:
|
||||
temporal = _temporal_value(value)
|
||||
if unit in {"year", "quarter", "month"}:
|
||||
months = amount * {"year": 12, "quarter": 3, "month": 1}[unit]
|
||||
return _add_months(temporal, months)
|
||||
seconds = amount * {
|
||||
"week": 7 * 24 * 60 * 60,
|
||||
"day": 24 * 60 * 60,
|
||||
"hour": 60 * 60,
|
||||
"minute": 60,
|
||||
"second": 1,
|
||||
}.get(unit, 0)
|
||||
if unit not in {"week", "day", "hour", "minute", "second"}:
|
||||
raise ValueError(f"Unsupported interval unit {unit!r}.")
|
||||
return temporal + timedelta(seconds=seconds)
|
||||
|
||||
|
||||
def _evaluate_date_diff(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> int:
|
||||
end = _temporal_value(_evaluate_child(expression, row))
|
||||
start = _temporal_value(_evaluate(expression.expression, row))
|
||||
unit = _date_unit(expression.args["unit"])
|
||||
if unit == "year":
|
||||
return end.year - start.year
|
||||
if unit in {"quarter", "month"}:
|
||||
months = (end.year - start.year) * 12 + end.month - start.month
|
||||
return months // 3 if unit == "quarter" else months
|
||||
delta = datetime.combine(end, datetime.min.time()) - datetime.combine(
|
||||
start,
|
||||
datetime.min.time(),
|
||||
) if not isinstance(end, datetime) and not isinstance(start, datetime) else (
|
||||
_as_datetime(end) - _as_datetime(start)
|
||||
)
|
||||
divisor = {
|
||||
"week": 7 * 24 * 60 * 60,
|
||||
"day": 24 * 60 * 60,
|
||||
"hour": 60 * 60,
|
||||
"minute": 60,
|
||||
"second": 1,
|
||||
}.get(unit)
|
||||
if divisor is None:
|
||||
raise ValueError(f"Unsupported date difference unit {unit!r}.")
|
||||
return int(delta.total_seconds() / divisor)
|
||||
|
||||
|
||||
def _interval_value(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> tuple[int, str]:
|
||||
if not isinstance(expression, exp.Interval):
|
||||
raise ValueError("Date arithmetic requires an INTERVAL.")
|
||||
interval = _evaluate_interval(expression, row)
|
||||
return interval.amount, interval.unit
|
||||
|
||||
|
||||
def _evaluate_interval(
|
||||
expression: exp.Expression,
|
||||
row: dict[str, Any],
|
||||
) -> _IntervalValue:
|
||||
return _IntervalValue(
|
||||
amount=int(_evaluate(expression.this, row)),
|
||||
unit=_date_unit(expression.args["unit"]),
|
||||
)
|
||||
|
||||
|
||||
def _date_unit(expression: exp.Expression) -> str:
|
||||
unit = str(expression.this).casefold()
|
||||
aliases = {
|
||||
"years": "year",
|
||||
"quarters": "quarter",
|
||||
"months": "month",
|
||||
"weeks": "week",
|
||||
"days": "day",
|
||||
"hours": "hour",
|
||||
"minutes": "minute",
|
||||
"seconds": "second",
|
||||
}
|
||||
return aliases.get(unit, unit)
|
||||
|
||||
|
||||
def _temporal_value(value: Any) -> date | datetime:
|
||||
if isinstance(value, datetime):
|
||||
return value
|
||||
if isinstance(value, date):
|
||||
return value
|
||||
text = str(value or "").strip()
|
||||
if not text:
|
||||
raise ValueError("A date or date-time value is required.")
|
||||
if "T" not in text and " " not in text:
|
||||
return date.fromisoformat(text)
|
||||
return datetime.fromisoformat(text.replace("Z", "+00:00"))
|
||||
|
||||
|
||||
def _as_datetime(value: date | datetime) -> datetime:
|
||||
return (
|
||||
value
|
||||
if isinstance(value, datetime)
|
||||
else datetime.combine(value, datetime.min.time())
|
||||
)
|
||||
|
||||
|
||||
def _add_months(value: date | datetime, months: int) -> date | datetime:
|
||||
month_index = value.year * 12 + value.month - 1 + months
|
||||
year, zero_based_month = divmod(month_index, 12)
|
||||
month = zero_based_month + 1
|
||||
day = min(value.day, calendar.monthrange(year, month)[1])
|
||||
return value.replace(year=year, month=month, day=day)
|
||||
|
||||
|
||||
def _evaluate_cast(expression: exp.Expression, row: dict[str, Any]) -> Any:
|
||||
return convert_value(
|
||||
_evaluate_child(expression, row),
|
||||
@@ -350,7 +915,6 @@ _COMPARISON_OPERATIONS: dict[
|
||||
_STRING_OPERATIONS: dict[type[exp.Expression], Callable[[str], str]] = {
|
||||
exp.Lower: str.lower,
|
||||
exp.Upper: str.upper,
|
||||
exp.Trim: str.strip,
|
||||
}
|
||||
_EVALUATORS: dict[
|
||||
type[exp.Expression],
|
||||
@@ -361,11 +925,18 @@ _EVALUATORS: dict[
|
||||
exp.Null: lambda _expression, _row: None,
|
||||
exp.Boolean: lambda expression, _row: bool(expression.this),
|
||||
exp.Literal: lambda expression, _row: _literal(expression), # type: ignore[arg-type]
|
||||
exp.Interval: _evaluate_interval,
|
||||
exp.Neg: _evaluate_negation,
|
||||
exp.Not: _evaluate_boolean_not,
|
||||
exp.And: _evaluate_boolean_binary,
|
||||
exp.Or: _evaluate_boolean_binary,
|
||||
exp.Is: _evaluate_is,
|
||||
exp.Like: _evaluate_like,
|
||||
exp.ILike: _evaluate_like,
|
||||
exp.In: _evaluate_in,
|
||||
exp.Between: _evaluate_between,
|
||||
exp.RegexpLike: _evaluate_regexp_match,
|
||||
exp.RegexpFullMatch: _evaluate_regexp_match,
|
||||
**{
|
||||
expression_type: _evaluate_arithmetic
|
||||
for expression_type in _ARITHMETIC_OPERATIONS
|
||||
@@ -379,12 +950,32 @@ _EVALUATORS: dict[
|
||||
for expression_type in _STRING_OPERATIONS
|
||||
},
|
||||
exp.Length: _evaluate_length,
|
||||
exp.Trim: _evaluate_trim,
|
||||
exp.Abs: _evaluate_abs,
|
||||
exp.Round: _evaluate_round,
|
||||
exp.Coalesce: _evaluate_coalesce,
|
||||
exp.Concat: _evaluate_concat,
|
||||
exp.ConcatWs: _evaluate_concat_ws,
|
||||
exp.DPipe: _evaluate_pipe,
|
||||
exp.Greatest: _evaluate_extreme,
|
||||
exp.Least: _evaluate_extreme,
|
||||
exp.Nullif: _evaluate_nullif,
|
||||
exp.Pad: _evaluate_pad,
|
||||
exp.Replace: _evaluate_replace,
|
||||
exp.RegexpReplace: _evaluate_regexp_replace,
|
||||
exp.SplitPart: _evaluate_split_part,
|
||||
exp.Substring: _evaluate_substring,
|
||||
exp.TimeToStr: _evaluate_time_to_str,
|
||||
exp.ToChar: _evaluate_to_char,
|
||||
exp.Translate: _evaluate_translate,
|
||||
exp.Extract: _evaluate_extract,
|
||||
exp.Year: _evaluate_temporal_part,
|
||||
exp.Month: _evaluate_temporal_part,
|
||||
exp.Day: _evaluate_temporal_part,
|
||||
exp.TimestampTrunc: _evaluate_timestamp_trunc,
|
||||
exp.DateAdd: _evaluate_date_add,
|
||||
exp.DateSub: _evaluate_date_add,
|
||||
exp.DateDiff: _evaluate_date_diff,
|
||||
exp.Cast: _evaluate_cast,
|
||||
exp.Case: _evaluate_case,
|
||||
exp.If: _evaluate_if,
|
||||
@@ -462,6 +1053,11 @@ def _infer_numeric(
|
||||
expression: exp.Expression,
|
||||
schema: dict[str, ExpressionDataType],
|
||||
) -> ExpressionDataType:
|
||||
if (
|
||||
isinstance(expression, (exp.Add, exp.Sub))
|
||||
and isinstance(expression.expression, exp.Interval)
|
||||
):
|
||||
return _infer(expression.this, schema)
|
||||
child_types = {
|
||||
_infer(item, schema)
|
||||
for item in expression.iter_expressions()
|
||||
@@ -485,6 +1081,18 @@ def _infer_coalesce(
|
||||
)
|
||||
|
||||
|
||||
def _infer_extreme(
|
||||
expression: exp.Expression,
|
||||
schema: dict[str, ExpressionDataType],
|
||||
) -> ExpressionDataType:
|
||||
return _common_expression_type(
|
||||
[
|
||||
_infer(item, schema)
|
||||
for item in (expression.this, *expression.expressions)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _infer_case(
|
||||
expression: exp.Expression,
|
||||
schema: dict[str, ExpressionDataType],
|
||||
@@ -553,6 +1161,12 @@ _TYPE_INFERERS: dict[
|
||||
exp.And,
|
||||
exp.Or,
|
||||
exp.Is,
|
||||
exp.Like,
|
||||
exp.ILike,
|
||||
exp.In,
|
||||
exp.Between,
|
||||
exp.RegexpLike,
|
||||
exp.RegexpFullMatch,
|
||||
exp.Not,
|
||||
)
|
||||
},
|
||||
@@ -564,8 +1178,16 @@ _TYPE_INFERERS: dict[
|
||||
exp.Upper,
|
||||
exp.Trim,
|
||||
exp.Concat,
|
||||
exp.ConcatWs,
|
||||
exp.DPipe,
|
||||
exp.Pad,
|
||||
exp.Replace,
|
||||
exp.RegexpReplace,
|
||||
exp.SplitPart,
|
||||
exp.Substring,
|
||||
exp.TimeToStr,
|
||||
exp.ToChar,
|
||||
exp.Translate,
|
||||
)
|
||||
},
|
||||
exp.Cast: _infer_cast,
|
||||
@@ -582,6 +1204,17 @@ _TYPE_INFERERS: dict[
|
||||
)
|
||||
},
|
||||
exp.Coalesce: _infer_coalesce,
|
||||
exp.Greatest: _infer_extreme,
|
||||
exp.Least: _infer_extreme,
|
||||
exp.Nullif: _infer_child,
|
||||
exp.Extract: lambda _expression, _schema: "integer",
|
||||
exp.Year: lambda _expression, _schema: "integer",
|
||||
exp.Month: lambda _expression, _schema: "integer",
|
||||
exp.Day: lambda _expression, _schema: "integer",
|
||||
exp.DateDiff: lambda _expression, _schema: "integer",
|
||||
exp.DateAdd: _infer_child,
|
||||
exp.DateSub: _infer_child,
|
||||
exp.TimestampTrunc: _infer_child,
|
||||
exp.Case: _infer_case,
|
||||
exp.If: _infer_if,
|
||||
exp.Paren: _infer_child,
|
||||
|
||||
@@ -47,7 +47,8 @@ DERIVE_OPERATIONS = frozenset(
|
||||
"divide",
|
||||
}
|
||||
)
|
||||
JOIN_TYPES = frozenset({"inner", "left", "right", "full"})
|
||||
JOIN_TYPES = frozenset({"inner", "left", "right", "full", "semi", "anti"})
|
||||
RANK_METHODS = frozenset({"row_number", "rank", "dense_rank"})
|
||||
DATA_TYPES = frozenset(
|
||||
{"string", "integer", "number", "boolean", "date", "datetime"}
|
||||
)
|
||||
@@ -901,6 +902,65 @@ def _validate_expression(node: GraphNode) -> list[DataflowDiagnostic]:
|
||||
return diagnostics
|
||||
|
||||
|
||||
def _validate_calculate(node: GraphNode) -> list[DataflowDiagnostic]:
|
||||
calculations = node.config.get("calculations")
|
||||
if not isinstance(calculations, list) or not 1 <= len(calculations) <= 100:
|
||||
return [
|
||||
_node_field_error(
|
||||
node,
|
||||
"calculate.required",
|
||||
"Add between one and 100 calculated columns.",
|
||||
"calculations",
|
||||
)
|
||||
]
|
||||
diagnostics: list[DataflowDiagnostic] = []
|
||||
targets: list[str] = []
|
||||
for item in calculations:
|
||||
if not isinstance(item, dict):
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"calculate.item",
|
||||
"Every calculation needs a target column and expression.",
|
||||
"calculations",
|
||||
)
|
||||
)
|
||||
continue
|
||||
target = str(item.get("target_column") or "").strip()
|
||||
expression = str(item.get("expression") or "").strip()
|
||||
result_type = str(item.get("result_type") or "unknown")
|
||||
if not target or not expression:
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"calculate.item",
|
||||
"Every calculation needs a target column and expression.",
|
||||
"calculations",
|
||||
)
|
||||
)
|
||||
if target:
|
||||
targets.append(target)
|
||||
if result_type not in {"unknown", *DATA_TYPES}:
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"calculate.result_type",
|
||||
"Choose a supported calculation result type.",
|
||||
"calculations",
|
||||
)
|
||||
)
|
||||
if len(targets) != len(set(targets)):
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"calculate.duplicate_target",
|
||||
"Calculated column names must be unique within the block.",
|
||||
"calculations",
|
||||
)
|
||||
)
|
||||
return diagnostics
|
||||
|
||||
|
||||
def _validate_convert(node: GraphNode) -> list[DataflowDiagnostic]:
|
||||
diagnostics = _validate_source_target_columns(node, "convert")
|
||||
if node.config.get("target_type") not in DATA_TYPES:
|
||||
@@ -988,6 +1048,56 @@ def _validate_sort(node: GraphNode) -> list[DataflowDiagnostic]:
|
||||
]
|
||||
|
||||
|
||||
def _validate_rank(node: GraphNode) -> list[DataflowDiagnostic]:
|
||||
diagnostics: list[DataflowDiagnostic] = []
|
||||
if node.config.get("method", "row_number") not in RANK_METHODS:
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"rank.method",
|
||||
"Choose row number, rank, or dense rank.",
|
||||
"method",
|
||||
)
|
||||
)
|
||||
if not _non_empty_text(node.config.get("target_column")):
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"rank.target_column",
|
||||
"Choose an output column.",
|
||||
"target_column",
|
||||
)
|
||||
)
|
||||
partition_by = node.config.get("partition_by", [])
|
||||
if not isinstance(partition_by, list) or any(
|
||||
not _non_empty_text(item) for item in partition_by
|
||||
):
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"rank.partition_by",
|
||||
"Partition columns must be named.",
|
||||
"partition_by",
|
||||
)
|
||||
)
|
||||
order_by = node.config.get("order_by")
|
||||
if not isinstance(order_by, list) or not order_by or not all(
|
||||
isinstance(item, dict)
|
||||
and _non_empty_text(item.get("column"))
|
||||
and item.get("direction", "asc") in {"asc", "desc"}
|
||||
for item in order_by
|
||||
):
|
||||
diagnostics.append(
|
||||
_node_field_error(
|
||||
node,
|
||||
"rank.order_by",
|
||||
"Add at least one ordered column with a valid direction.",
|
||||
"order_by",
|
||||
)
|
||||
)
|
||||
return diagnostics
|
||||
|
||||
|
||||
def _validate_limit(node: GraphNode) -> list[DataflowDiagnostic]:
|
||||
count = node.config.get("count")
|
||||
if (
|
||||
@@ -1296,10 +1406,12 @@ def _register_config_validators() -> None:
|
||||
"select": _validate_select,
|
||||
"derive": _validate_derive,
|
||||
"expression": _validate_expression,
|
||||
"calculate": _validate_calculate,
|
||||
"convert": _validate_convert,
|
||||
"replace": _validate_replace,
|
||||
"aggregate": _validate_aggregate,
|
||||
"sort": _validate_sort,
|
||||
"window.rank": _validate_rank,
|
||||
"limit": _validate_limit,
|
||||
"quality.rules": _validate_quality,
|
||||
"reconcile.compare": _validate_reconcile,
|
||||
@@ -1323,6 +1435,7 @@ __all__ = [
|
||||
"DERIVE_OPERATIONS",
|
||||
"FILTER_OPERATORS",
|
||||
"JOIN_TYPES",
|
||||
"RANK_METHODS",
|
||||
"SUPPORTED_NODE_TYPES",
|
||||
"canonical_graph_payload",
|
||||
"definition_hash",
|
||||
|
||||
@@ -279,26 +279,43 @@ def _node_expressions(
|
||||
node: GraphNode,
|
||||
output_state: SchemaState,
|
||||
) -> tuple[IrExpression, ...]:
|
||||
if node.type == "calculate":
|
||||
sources = [
|
||||
str(item.get("expression") or "")
|
||||
for item in node.config.get("calculations", [])
|
||||
if isinstance(item, dict) and item.get("expression")
|
||||
]
|
||||
return tuple(
|
||||
item
|
||||
for source in sources
|
||||
if (item := _ir_expression(source, output_state)) is not None
|
||||
)
|
||||
source = node.config.get("expression")
|
||||
if node.type not in {"expression", "filter.expression"} or not source:
|
||||
return ()
|
||||
expression = _ir_expression(str(source), output_state)
|
||||
return (expression,) if expression is not None else ()
|
||||
|
||||
|
||||
def _ir_expression(
|
||||
source: str,
|
||||
output_state: SchemaState,
|
||||
) -> IrExpression | None:
|
||||
try:
|
||||
parsed = parse_expression(str(source))
|
||||
except ExpressionError:
|
||||
return ()
|
||||
return None
|
||||
inferred = infer_expression_type(parsed, output_state.types)
|
||||
return (
|
||||
IrExpression(
|
||||
source=parsed.source,
|
||||
columns=parsed.columns,
|
||||
result_type=_data_type(inferred),
|
||||
semantic_hash=_hash_payload(
|
||||
{
|
||||
"dialect": "duckdb",
|
||||
"source": parsed.sql(),
|
||||
"result_type": inferred,
|
||||
}
|
||||
),
|
||||
return IrExpression(
|
||||
source=parsed.source,
|
||||
columns=parsed.columns,
|
||||
result_type=_data_type(inferred),
|
||||
semantic_hash=_hash_payload(
|
||||
{
|
||||
"dialect": "duckdb",
|
||||
"source": parsed.sql(),
|
||||
"result_type": inferred,
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -117,6 +117,8 @@ _NODE_TYPES = (
|
||||
("left", "All left rows"),
|
||||
("right", "All right rows"),
|
||||
("full", "All rows"),
|
||||
("semi", "Left rows with a match"),
|
||||
("anti", "Left rows without a match"),
|
||||
),
|
||||
),
|
||||
NodeConfigField(id="left_keys", label="Left keys", kind="column_list", required=True),
|
||||
@@ -289,6 +291,35 @@ _NODE_TYPES = (
|
||||
},
|
||||
sql_support="partial",
|
||||
),
|
||||
NodeTypeDefinition(
|
||||
type="calculate",
|
||||
category="transform",
|
||||
label="Calculate columns",
|
||||
description=(
|
||||
"Create or replace several columns using ordered safe typed "
|
||||
"expressions."
|
||||
),
|
||||
icon="calculator",
|
||||
input_ports=(NodePortDefinition(id="input", label="Input"),),
|
||||
config_fields=(
|
||||
NodeConfigField(
|
||||
id="calculations",
|
||||
label="Calculations",
|
||||
kind="calculations",
|
||||
required=True,
|
||||
),
|
||||
),
|
||||
default_config={
|
||||
"calculations": [
|
||||
{
|
||||
"target_column": "",
|
||||
"expression": "",
|
||||
"result_type": "unknown",
|
||||
}
|
||||
]
|
||||
},
|
||||
sql_support="partial",
|
||||
),
|
||||
NodeTypeDefinition(
|
||||
type="convert",
|
||||
category="transform",
|
||||
@@ -386,6 +417,53 @@ _NODE_TYPES = (
|
||||
config_fields=(NodeConfigField(id="fields", label="Sort fields", kind="sort_fields", required=True),),
|
||||
default_config={"fields": [{"column": "", "direction": "asc"}]},
|
||||
),
|
||||
NodeTypeDefinition(
|
||||
type="window.rank",
|
||||
category="transform",
|
||||
label="Rank rows",
|
||||
description=(
|
||||
"Number or rank rows within optional partitions using a stable "
|
||||
"ordering."
|
||||
),
|
||||
icon="list-ordered",
|
||||
input_ports=(NodePortDefinition(id="input", label="Input"),),
|
||||
config_fields=(
|
||||
NodeConfigField(
|
||||
id="method",
|
||||
label="Method",
|
||||
kind="select",
|
||||
options=(
|
||||
("row_number", "Row number"),
|
||||
("rank", "Rank with gaps"),
|
||||
("dense_rank", "Dense rank"),
|
||||
),
|
||||
),
|
||||
NodeConfigField(
|
||||
id="target_column",
|
||||
label="Output column",
|
||||
kind="text",
|
||||
required=True,
|
||||
),
|
||||
NodeConfigField(
|
||||
id="partition_by",
|
||||
label="Partition by",
|
||||
kind="column_list",
|
||||
),
|
||||
NodeConfigField(
|
||||
id="order_by",
|
||||
label="Order by",
|
||||
kind="sort_fields",
|
||||
required=True,
|
||||
),
|
||||
),
|
||||
default_config={
|
||||
"method": "row_number",
|
||||
"target_column": "row_number",
|
||||
"partition_by": [],
|
||||
"order_by": [{"column": "", "direction": "asc"}],
|
||||
},
|
||||
sql_support="partial",
|
||||
),
|
||||
NodeTypeDefinition(
|
||||
type="limit",
|
||||
category="transform",
|
||||
|
||||
@@ -185,6 +185,8 @@ def _join(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
||||
field_name="right_keys",
|
||||
),
|
||||
]
|
||||
if node.config.get("join_type", "inner") in {"semi", "anti"}:
|
||||
return SchemaPropagationResult(left_state, tuple(diagnostics))
|
||||
prefix = str(node.config.get("right_prefix", "right_"))
|
||||
prefixed_right = {
|
||||
f"{prefix}{column}"
|
||||
@@ -377,6 +379,66 @@ def _expression(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
||||
return SchemaPropagationResult(state, tuple(diagnostics))
|
||||
|
||||
|
||||
def _calculate(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
||||
node = context.node
|
||||
state = context.input_state
|
||||
diagnostics: list[DataflowDiagnostic] = []
|
||||
for item in _mapping_items(node.config.get("calculations")):
|
||||
source = str(item.get("expression") or "")
|
||||
target = str(item.get("target_column") or "")
|
||||
try:
|
||||
parsed = parse_expression(source)
|
||||
except ExpressionError as exc:
|
||||
diagnostics.append(
|
||||
_error(
|
||||
"expression.invalid",
|
||||
str(exc),
|
||||
node_id=node.id,
|
||||
field="calculations",
|
||||
)
|
||||
)
|
||||
continue
|
||||
diagnostics.extend(
|
||||
_unknown_columns(
|
||||
node,
|
||||
state,
|
||||
list(parsed.columns),
|
||||
field_name="calculations",
|
||||
)
|
||||
)
|
||||
inferred = infer_expression_type(
|
||||
parsed,
|
||||
{
|
||||
name: state.type_of(name) # type: ignore[dict-item]
|
||||
for name in state.columns
|
||||
},
|
||||
)
|
||||
expected = str(item.get("result_type") or "unknown")
|
||||
if expected != "unknown" and inferred not in {
|
||||
"unknown",
|
||||
"null",
|
||||
expected,
|
||||
}:
|
||||
diagnostics.append(
|
||||
_warning(
|
||||
"calculate.type_mismatch",
|
||||
(
|
||||
f"Calculation for {target!r} infers {inferred}, "
|
||||
f"not {expected}."
|
||||
),
|
||||
node_id=node.id,
|
||||
field="calculations",
|
||||
)
|
||||
)
|
||||
if target:
|
||||
state = _with_column(
|
||||
state,
|
||||
target,
|
||||
expected if expected != "unknown" else inferred,
|
||||
)
|
||||
return SchemaPropagationResult(state, tuple(diagnostics))
|
||||
|
||||
|
||||
def _convert_or_replace(
|
||||
context: SchemaPropagationContext,
|
||||
) -> SchemaPropagationResult:
|
||||
@@ -463,6 +525,31 @@ def _sort(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
||||
)
|
||||
|
||||
|
||||
def _rank(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
||||
node = context.node
|
||||
columns = [
|
||||
*_text_items(node.config.get("partition_by")),
|
||||
*[
|
||||
str(item.get("column"))
|
||||
for item in _mapping_items(node.config.get("order_by"))
|
||||
if item.get("column")
|
||||
],
|
||||
]
|
||||
diagnostics = _unknown_columns(
|
||||
node,
|
||||
context.input_state,
|
||||
columns,
|
||||
field_name="order_by",
|
||||
)
|
||||
target = str(node.config.get("target_column") or "")
|
||||
state = (
|
||||
_with_column(context.input_state, target, "integer")
|
||||
if target
|
||||
else context.input_state
|
||||
)
|
||||
return SchemaPropagationResult(state, tuple(diagnostics))
|
||||
|
||||
|
||||
def _quality(context: SchemaPropagationContext) -> SchemaPropagationResult:
|
||||
rules = _mapping_items(context.node.config.get("rules"))
|
||||
diagnostics = _unknown_columns(
|
||||
@@ -874,10 +961,12 @@ def register_schema_propagators() -> None:
|
||||
"select": _select,
|
||||
"derive": _derive,
|
||||
"expression": _expression,
|
||||
"calculate": _calculate,
|
||||
"convert": _convert_or_replace,
|
||||
"replace": _convert_or_replace,
|
||||
"aggregate": _aggregate,
|
||||
"sort": _sort,
|
||||
"window.rank": _rank,
|
||||
"limit": _identity,
|
||||
"quality.rules": _quality,
|
||||
"reconcile.compare": _reconcile,
|
||||
|
||||
@@ -1128,11 +1128,37 @@ def _join_config(
|
||||
) -> dict[str, Any]:
|
||||
kind = str(join.args.get("kind") or "").casefold()
|
||||
side = str(join.args.get("side") or "").casefold()
|
||||
if kind not in {"", "inner", "outer"} or side not in {"", "left", "right", "full"}:
|
||||
if kind not in {"", "inner", "outer", "semi", "anti"} or side not in {
|
||||
"",
|
||||
"left",
|
||||
"right",
|
||||
"full",
|
||||
}:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.join_type", "JOIN supports INNER, LEFT, RIGHT, or FULL joins only.")]
|
||||
[
|
||||
_sql_error(
|
||||
"sql.join_type",
|
||||
(
|
||||
"JOIN supports INNER, LEFT, RIGHT, FULL, SEMI, "
|
||||
"or ANTI joins only."
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
join_type = side or ("inner" if kind in {"", "inner"} else "")
|
||||
if kind in {"semi", "anti"} and side:
|
||||
raise SqlCompilationError(
|
||||
[
|
||||
_sql_error(
|
||||
"sql.join_type",
|
||||
"SEMI and ANTI joins cannot use a side qualifier.",
|
||||
)
|
||||
]
|
||||
)
|
||||
join_type = (
|
||||
kind
|
||||
if kind in {"semi", "anti"}
|
||||
else side or ("inner" if kind in {"", "inner"} else "")
|
||||
)
|
||||
if not join_type:
|
||||
raise SqlCompilationError(
|
||||
[_sql_error("sql.join_type", "OUTER JOIN requires LEFT, RIGHT, or FULL.")]
|
||||
@@ -1471,6 +1497,41 @@ def _render_expression(
|
||||
state.selected = True
|
||||
|
||||
|
||||
def _render_calculate(
|
||||
node: GraphNode,
|
||||
state: _SqlRenderState,
|
||||
column_expression: Callable[[str], exp.Column],
|
||||
) -> None:
|
||||
_require_projection_slot(node, state)
|
||||
calculated_targets: set[str] = set()
|
||||
expressions: list[exp.Expression] = [exp.Star()]
|
||||
for item in node.config["calculations"]:
|
||||
parsed = parse_expression(str(item["expression"]))
|
||||
dependencies = set(parsed.columns) & calculated_targets
|
||||
if dependencies:
|
||||
raise SqlCompilationError(
|
||||
[
|
||||
_node_sql_error(
|
||||
node.id,
|
||||
"sql.sequential_calculation",
|
||||
(
|
||||
"SQL view cannot render a calculated column that "
|
||||
"depends on an earlier calculation in the same block."
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
expression = _qualified_expression(
|
||||
str(item["expression"]),
|
||||
column_expression=column_expression,
|
||||
)
|
||||
target = str(item["target_column"])
|
||||
expressions.append(expression.as_(target))
|
||||
calculated_targets.add(target)
|
||||
state.select_expressions = expressions
|
||||
state.selected = True
|
||||
|
||||
|
||||
def _render_derive(
|
||||
node: GraphNode,
|
||||
state: _SqlRenderState,
|
||||
@@ -1600,6 +1661,41 @@ def _render_sort(
|
||||
]
|
||||
|
||||
|
||||
def _render_rank(
|
||||
node: GraphNode,
|
||||
state: _SqlRenderState,
|
||||
column_expression: Callable[[str], exp.Column],
|
||||
) -> None:
|
||||
_require_projection_slot(node, state)
|
||||
function = {
|
||||
"row_number": exp.RowNumber,
|
||||
"rank": exp.Rank,
|
||||
"dense_rank": exp.DenseRank,
|
||||
}[str(node.config.get("method", "row_number"))]()
|
||||
window = exp.Window(
|
||||
this=function,
|
||||
partition_by=[
|
||||
column_expression(str(column))
|
||||
for column in node.config.get("partition_by", [])
|
||||
],
|
||||
order=exp.Order(
|
||||
expressions=[
|
||||
exp.Ordered(
|
||||
this=column_expression(str(item["column"])),
|
||||
desc=item.get("direction", "asc") == "desc",
|
||||
nulls_first=False,
|
||||
)
|
||||
for item in node.config["order_by"]
|
||||
]
|
||||
),
|
||||
)
|
||||
state.select_expressions = [
|
||||
exp.Star(),
|
||||
window.as_(str(node.config["target_column"])),
|
||||
]
|
||||
state.selected = True
|
||||
|
||||
|
||||
def _render_limit(
|
||||
node: GraphNode,
|
||||
state: _SqlRenderState,
|
||||
@@ -1642,7 +1738,10 @@ def _require_projection_slot(
|
||||
_node_sql_error(
|
||||
node.id,
|
||||
"sql.multiple_select",
|
||||
"Only one expression, conversion, derive, select, or aggregate transform is supported in SQL view.",
|
||||
(
|
||||
"Only one calculation, ranking, conversion, derive, "
|
||||
"select, or aggregate transform is supported in SQL view."
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
@@ -1660,10 +1759,12 @@ def _register_sql_renderers() -> None:
|
||||
"select": _render_select,
|
||||
"derive": _render_derive,
|
||||
"expression": _render_expression,
|
||||
"calculate": _render_calculate,
|
||||
"convert": _render_convert,
|
||||
"replace": _render_unsupported,
|
||||
"aggregate": _render_aggregate,
|
||||
"sort": _render_sort,
|
||||
"window.rank": _render_rank,
|
||||
"limit": _render_limit,
|
||||
"quality.rules": _render_unsupported,
|
||||
"reconcile.compare": _render_unsupported,
|
||||
|
||||
Reference in New Issue
Block a user