feat(dataflow): model production SQL flow patterns

This commit is contained in:
2026-07-30 18:19:18 +02:00
parent 0664570607
commit e091042948
26 changed files with 2169 additions and 36 deletions
+114
View File
@@ -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"])]
),
+634 -1
View File
@@ -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,
+114 -1
View File
@@ -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",
+30 -13
View File
@@ -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,
+105 -4
View File
@@ -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,