From 5846842ac7fb649d2468d4f74f010e7be9df4ad1 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 22:12:54 +0000 Subject: [PATCH 01/15] tests: Spark 4.2.0 observe goldens (classic and Connect) and their generator (P3-T084) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../fixtures/spark-4.2.0/observe-goldens.json | 1833 +++++++++++++++++ tests/scripts/gen_observe_goldens.py | 693 +++++++ tests/scripts/merge_observe_goldens.py | 28 + 3 files changed, 2554 insertions(+) create mode 100644 tests/fixtures/spark-4.2.0/observe-goldens.json create mode 100644 tests/scripts/gen_observe_goldens.py create mode 100644 tests/scripts/merge_observe_goldens.py diff --git a/tests/fixtures/spark-4.2.0/observe-goldens.json b/tests/fixtures/spark-4.2.0/observe-goldens.json new file mode 100644 index 000000000..4c5e0adc0 --- /dev/null +++ b/tests/fixtures/spark-4.2.0/observe-goldens.json @@ -0,0 +1,1833 @@ +{ + "classic": { + "aggregate_with_alias_inside_nested_alias": { + "action_result": 3, + "get": "ok", + "metrics": { + "b": 3 + }, + "obs_name_given": true, + "types": { + "b": "int" + } + }, + "anonymous_observation": { + "get": "ok", + "metrics": { + "n": 3 + }, + "name_is_str": false, + "name_len": null, + "types": { + "n": "int" + } + }, + "ansi_off_overflow_sum": { + "action_result": 2, + "get": "ok", + "metrics": { + "p1": -9223372036854775808, + "s": -9223372036854775808 + }, + "obs_name_given": true, + "types": { + "p1": "int", + "s": "int" + } + }, + "ansi_on_cast_error_inside_metric": { + "errorClass": "CAST_INVALID_INPUT", + "message": "[CAST_INVALID_INPUT] The value 'a' of the type \"STRING\" cannot be cast to \"INT\" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. SQLSTATE: 22018", + "messageParameters": { + "ansiConfig": "\"spark.sql.ansi.enabled\"", + "expression": "'a'", + "sourceType": "\"STRING\"", + "targetType": "\"INT\"" + }, + "outcome": "error", + "python_class": "NumberFormatException", + "sqlState": "22018" + }, + "ansi_on_overflow_sum": { + "action_result": 2, + "get": "ok", + "metrics": { + "s": 9223372036854775807 + }, + "obs_name_given": true, + "types": { + "s": "int" + } + }, + "basic_alias": { + "action_result": 3, + "get": "ok", + "metrics": { + "mk": "c", + "rows": 3, + "sv": 3 + }, + "obs_name_given": true, + "types": { + "mk": "str", + "rows": "int", + "sv": "int" + } + }, + "cache": { + "after_first": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "after_second": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "counts": [ + 3, + 3 + ] + }, + "cache_then_new_observe_on_cached": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "composite_expr": { + "action_result": 3, + "get": "ok", + "metrics": { + "e": "4.25", + "n": 3 + }, + "obs_name_given": true, + "types": { + "e": "float", + "n": "int" + } + }, + "conditional_metric_no_raise_unselected": { + "action_result": 2, + "get": "ok", + "metrics": { + "s": "5.0" + }, + "obs_name_given": true, + "types": { + "s": "float" + } + }, + "count_action_prunes_columns": { + "action_result": 3, + "get": "ok", + "metrics": { + "s": 3 + }, + "obs_name_given": true, + "types": { + "s": "int" + } + }, + "empty_input": { + "action_result": 0, + "get": "ok", + "metrics": { + "a": null, + "c": 0, + "l": [], + "mx": null, + "s": null + }, + "obs_name_given": true, + "types": { + "a": "NoneType", + "c": "int", + "l": "list", + "mx": "NoneType", + "s": "NoneType" + } + }, + "empty_input_no_rows_collected": { + "action_result": [], + "get": "ok", + "metrics": { + "c": 0 + }, + "obs_name_given": true, + "types": { + "c": "int" + } + }, + "err_attribute": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE] Invalid observed metrics. Attribute \"v\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"v\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_attribute_alias_mixed": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE] Invalid observed metrics. Attribute \"v AS raw\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"v AS raw\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_distinct": { + "errorClass": "INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED] Invalid observed metrics. Aggregate expressions with DISTINCT are not allowed in observed metrics, but found: \"count(DISTINCT v) AS d\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"count(DISTINCT v) AS d\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_distinct_sum": { + "errorClass": "INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED] Invalid observed metrics. Aggregate expressions with DISTINCT are not allowed in observed metrics, but found: \"sum(DISTINCT v) AS d\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"sum(DISTINCT v) AS d\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_empty_name_string_api": { + "errorClass": "INVALID_OBSERVED_METRICS.MISSING_NAME", + "message": "[INVALID_OBSERVED_METRICS.MISSING_NAME] Invalid observed metrics. The observed metrics should be named: 'CollectMetrics , [count(1) AS n#], \n+- LogicalRDD [k#, v#, x#], false\n. SQLSTATE: 42K0E", + "messageParameters": { + "operator": "'CollectMetrics , [count(1) AS n#], \n+- LogicalRDD [k#, v#, x#], false\n" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_empty_observation_name": { + "errorClass": "VALUE_NOT_NON_EMPTY_STR", + "message": "[VALUE_NOT_NON_EMPTY_STR] Value for `name` must be a non-empty string, got ''.", + "messageParameters": { + "arg_name": "name", + "arg_value": "" + }, + "outcome": "error", + "python_class": "PySparkValueError", + "sqlState": null + }, + "err_filter": { + "errorClass": "INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_FILTER_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_FILTER_UNSUPPORTED] Invalid observed metrics. Aggregate expression with FILTER predicate are not allowed in observed metrics, but found: \"count(v) FILTER (WHERE (v > 1)) AS f\". SQLSTATE: 42K0E; line 1 pos 0", + "messageParameters": { + "expr": "\"count(v) FILTER (WHERE (v > 1)) AS f\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_first_ignorenulls_ok": { + "get": "ok", + "metrics": { + "f": 1 + }, + "outcome": "ok_unexpected", + "types": { + "f": "int" + } + }, + "err_generator": { + "errorClass": "UNSUPPORTED_GENERATOR.OUTSIDE_SELECT", + "message": "[UNSUPPORTED_GENERATOR.OUTSIDE_SELECT] The generator is not supported: outside the SELECT clause, found: 'CollectMetrics m, [explode(array(1, 2)) AS g#], . SQLSTATE: 42K0E", + "messageParameters": {}, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_get_before_observe": { + "errorClass": "NO_OBSERVE_BEFORE_GET", + "get": "error", + "message": "[NO_OBSERVE_BEFORE_GET] Should observe by calling `DataFrame.observe` before `get`.", + "messageParameters": {}, + "python_class": "PySparkAssertionError", + "sqlState": null + }, + "err_nested_aggregate": { + "errorClass": "INVALID_OBSERVED_METRICS.NESTED_AGGREGATES_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.NESTED_AGGREGATES_UNSUPPORTED] Invalid observed metrics. Nested aggregates are not allowed in observed metrics, but found: \"sum(max(v)) AS n\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"sum(max(v)) AS n\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_no_exprs": { + "errorClass": "CANNOT_BE_EMPTY", + "message": "[CANNOT_BE_EMPTY] At least one exprs must be specified.", + "messageParameters": { + "item": "exprs" + }, + "outcome": "error", + "python_class": "PySparkValueError", + "sqlState": null + }, + "err_no_exprs_string": { + "errorClass": "CANNOT_BE_EMPTY", + "message": "[CANNOT_BE_EMPTY] At least one exprs must be specified.", + "messageParameters": { + "item": "exprs" + }, + "outcome": "error", + "python_class": "PySparkValueError", + "sqlState": null + }, + "err_nondeterministic_plus_agg": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC] Invalid observed metrics. Non-deterministic expression \"(count(1) + rand()) AS r\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"(count(1) + rand()) AS r\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_nondeterministic_rand": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC] Invalid observed metrics. Non-deterministic expression \"rand() AS r\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"rand() AS r\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_observation_name_type": { + "errorClass": "NOT_EXPECTED_TYPE", + "message": "[NOT_EXPECTED_TYPE] Argument `name` should be str, got int.", + "messageParameters": { + "arg_name": "name", + "arg_type": "int", + "expected_type": "str" + }, + "outcome": "error", + "python_class": "PySparkTypeError", + "sqlState": "42K09" + }, + "err_reuse": { + "errorClass": "REUSE_OBSERVATION", + "message": "[REUSE_OBSERVATION] An Observation can be used with a DataFrame only once.", + "messageParameters": {}, + "outcome": "error", + "python_class": "PySparkAssertionError", + "sqlState": null + }, + "err_scalar_subquery": { + "get": "ok", + "metrics": { + "s": 2 + }, + "outcome": "ok_unexpected", + "types": { + "s": "int" + } + }, + "err_unresolved_column": { + "errorClass": "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "message": "[UNRESOLVED_COLUMN.WITH_SUGGESTION] A column, variable, or function parameter with name `nope` cannot be resolved. Did you mean one of the following? [`k`, `v`, `x`]. SQLSTATE: 42703", + "messageParameters": { + "objectName": "`nope`", + "proposal": "`k`, `v`, `x`" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42703" + }, + "err_unresolved_star": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE] Invalid observed metrics. Attribute \"k\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"k\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_window": { + "errorClass": "INVALID_OBSERVED_METRICS.WINDOW_EXPRESSIONS_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.WINDOW_EXPRESSIONS_UNSUPPORTED] Invalid observed metrics. Window expressions are not allowed in observed metrics, but found: \"sum(v) OVER (ORDER BY k ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS w\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"sum(v) OVER (ORDER BY k ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS w\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "exact_under_parallelism": { + "action_result": 100000, + "get": "ok", + "metrics": { + "mx": 99999, + "n": 100000, + "s": 4999950000 + }, + "obs_name_given": true, + "types": { + "mx": "int", + "n": "int", + "s": "int" + } + }, + "explain_does_not_trigger": { + "get": "blocked" + }, + "failed_query": { + "action": { + "errorClass": "DIVIDE_BY_ZERO", + "message": "[DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error. SQLSTATE: 22012", + "messageParameters": { + "config": "\"spark.sql.ansi.enabled\"" + }, + "python_class": "ArithmeticException", + "sqlState": "22012" + }, + "get": "ok", + "metrics": "", + "types": "" + }, + "failed_query_then_reuse_df_succeeds": { + "after_fail": { + "get": "ok", + "metrics": "", + "types": "" + }, + "after_success": { + "get": "ok", + "metrics": "", + "types": "" + } + }, + "filter_after_observe": { + "action_result": 1, + "get": "ok", + "metrics": { + "n": 3 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "filter_before_observe": { + "get": "ok", + "metrics": { + "n": 1 + }, + "types": { + "n": "int" + } + }, + "first_then_second_action": { + "after_count": { + "n": 1 + }, + "after_limit1": { + "n": 1 + } + }, + "get_after_transformation_only_no_action_then_action": { + "after_action": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "before_action": { + "get": "blocked" + } + }, + "get_blocks_without_action": { + "get": "blocked" + }, + "join_two_observed": { + "a": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "b": { + "get": "ok", + "metrics": { + "s": 3 + }, + "types": { + "s": "int" + } + }, + "count": 3 + }, + "limit_after_observe": { + "action_result": [ + [ + "a", + 1, + "1.5" + ] + ], + "get": "ok", + "metrics": { + "n": 1 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "limit_big_range": { + "action_result": [ + [ + 0 + ], + [ + 1 + ], + [ + 2 + ] + ], + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "obs_name_given": true, + "types": { + "n": "int", + "s": "int" + } + }, + "limit_zero": { + "action_result": [], + "get": "error", + "message": "java.lang.AssertionError in PythonSQLUtils.toPyRow", + "obs_name_given": true, + "python_class": "Py4JJavaError" + }, + "literal_only": { + "action_result": 3, + "get": "ok", + "metrics": { + "c": 42, + "s": "s" + }, + "obs_name_given": true, + "types": { + "c": "int", + "s": "str" + } + }, + "multiple_actions_first_wins": { + "after_first": { + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "types": { + "n": "int", + "s": "int" + } + }, + "after_more": { + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "types": { + "n": "int", + "s": "int" + } + }, + "counts": [ + 3, + 2, + 3 + ] + }, + "nested_observe_different_names": { + "a": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "b": { + "get": "ok", + "metrics": { + "m": 1 + }, + "types": { + "m": "int" + } + }, + "count": 1 + }, + "nested_observe_same_name_diff_df": { + "errorClass": "DUPLICATED_METRICS_NAME", + "message": "[DUPLICATED_METRICS_NAME] The metric name is not unique: dup. The same name cannot be used for metrics with different results.\nHowever multiple instances of metrics with with same result and name are allowed (e.g. self-joins). SQLSTATE: 42710", + "messageParameters": { + "metricName": "dup" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42710" + }, + "no_alias_names": { + "action_result": 3, + "get": "ok", + "metrics": { + "42": 42, + "count(1)": 3, + "max(v)": 2, + "sum((v + 1))": 5 + }, + "obs_name_given": true, + "types": { + "42": "int", + "count(1)": "int", + "max(v)": "int", + "sum((v + 1))": "int" + } + }, + "nulls_all": { + "action_result": 1, + "get": "ok", + "metrics": { + "c": 0, + "m": null, + "s": null + }, + "obs_name_given": true, + "types": { + "c": "int", + "m": "NoneType", + "s": "NoneType" + } + }, + "observation_in_groupby_downstream": { + "get": "ok", + "metrics": { + "n": 3 + }, + "rows": [ + [ + "a", + 1 + ], + [ + "b", + 1 + ], + [ + "c", + 1 + ] + ], + "types": { + "n": "int" + } + }, + "observation_on_action_collect_schema_only": { + "columns": [ + "k", + "v", + "x" + ], + "get_after_schema": { + "get": "blocked" + }, + "schema": "struct" + }, + "observation_through_sort_limit_topn": { + "action_result": [ + [ + 0 + ], + [ + 1 + ], + [ + 2 + ], + [ + 3 + ], + [ + 4 + ] + ], + "get": "ok", + "metrics": { + "n": 20 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "observed_metric_column_name_collision": { + "action_result": 3, + "get": "ok", + "metrics": { + "n": 3 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "ok_aggregate_over_subquery_free_udf_like": { + "get": "ok", + "metrics": { + "s": 3 + }, + "outcome": "ok_unexpected", + "types": { + "s": "int" + } + }, + "ok_approx_count_distinct": { + "get": "ok", + "metrics": { + "a": 2 + }, + "outcome": "ok_unexpected", + "types": { + "a": "int" + } + }, + "ok_count_if": { + "get": "ok", + "metrics": { + "c": 2 + }, + "outcome": "ok_unexpected", + "types": { + "c": "int" + } + }, + "ok_nondeterministic_inside_agg": { + "get": "ok", + "metrics": "", + "outcome": "ok_unexpected", + "types": "" + }, + "python_attrs": { + "has_future": false, + "has_get_async": false, + "has_name": false, + "public": [ + "_jo", + "_jvm", + "_name", + "_on", + "get" + ] + }, + "same_name_string_same_df_in_both_branches": { + "count": 6 + }, + "self_join_same_df": { + "count": 3, + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "self_union_same_df": { + "count": 6, + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "show_action_limit20": { + "action_result": 0, + "get": "ok", + "metrics": { + "n": 4 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "string_observe_returns_df_only": { + "count": 3, + "type": "DataFrame" + }, + "take_after_observe": { + "action_result": [ + [ + "a", + 1, + "1.5" + ] + ], + "get": "ok", + "metrics": { + "n": 1 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "types": { + "action_result": 1, + "get": "ok", + "metrics": { + "ai": "1.0", + "l": [ + 1 + ], + "mn": "nan", + "mz": "0.0", + "sd": { + "py_type": "Decimal", + "repr": "1.10" + }, + "si": 1, + "st": [ + 1, + 1 + ] + }, + "obs_name_given": true, + "types": { + "ai": "float", + "l": "list", + "mn": "float", + "mz": "float", + "sd": "Decimal", + "si": "int", + "st": "Row" + } + }, + "union_two_branches_distinct_names": { + "a": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "b": { + "get": "ok", + "metrics": { + "n": 1 + }, + "types": { + "n": "int" + } + }, + "count": 4 + }, + "union_two_branches_same_name_diff_df": { + "errorClass": "DUPLICATED_METRICS_NAME", + "message": "[DUPLICATED_METRICS_NAME] The metric name is not unique: dup. The same name cannot be used for metrics with different results.\nHowever multiple instances of metrics with with same result and name are allowed (e.g. self-joins). SQLSTATE: 42710", + "messageParameters": { + "metricName": "dup" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42710" + }, + "write_action": { + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "types": { + "n": "int", + "s": "int" + } + } + }, + "connect": { + "aggregate_with_alias_inside_nested_alias": { + "action_result": 3, + "get": "ok", + "metrics": { + "b": 3 + }, + "obs_name_given": true, + "types": { + "b": "int" + } + }, + "anonymous_observation": { + "get": "ok", + "metrics": { + "n": 3 + }, + "name_is_str": true, + "name_len": 36, + "types": { + "n": "int" + } + }, + "ansi_off_overflow_sum": { + "action_result": 2, + "get": "ok", + "metrics": { + "p1": -9223372036854775808, + "s": -9223372036854775808 + }, + "obs_name_given": true, + "types": { + "p1": "int", + "s": "int" + } + }, + "ansi_on_cast_error_inside_metric": { + "errorClass": "CAST_INVALID_INPUT", + "message": "[CAST_INVALID_INPUT] The value 'a' of the type \"STRING\" cannot be cast to \"INT\" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. SQLSTATE: 22018", + "messageParameters": { + "ansiConfig": "\"spark.sql.ansi.enabled\"", + "expression": "'a'", + "sourceType": "\"STRING\"", + "targetType": "\"INT\"" + }, + "outcome": "error", + "python_class": "NumberFormatException", + "sqlState": "22018" + }, + "ansi_on_overflow_sum": { + "action_result": 2, + "get": "ok", + "metrics": { + "s": 9223372036854775807 + }, + "obs_name_given": true, + "types": { + "s": "int" + } + }, + "basic_alias": { + "action_result": 3, + "get": "ok", + "metrics": { + "mk": "c", + "rows": 3, + "sv": 3 + }, + "obs_name_given": true, + "types": { + "mk": "str", + "rows": "int", + "sv": "int" + } + }, + "cache": { + "after_first": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "after_second": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "counts": [ + 3, + 3 + ] + }, + "cache_then_new_observe_on_cached": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "composite_expr": { + "action_result": 3, + "get": "ok", + "metrics": { + "e": "4.25", + "n": 3 + }, + "obs_name_given": true, + "types": { + "e": "float", + "n": "int" + } + }, + "conditional_metric_no_raise_unselected": { + "action_result": 2, + "get": "ok", + "metrics": { + "s": "5.0" + }, + "obs_name_given": true, + "types": { + "s": "float" + } + }, + "count_action_prunes_columns": { + "action_result": 3, + "get": "ok", + "metrics": { + "s": 3 + }, + "obs_name_given": true, + "types": { + "s": "int" + } + }, + "empty_input": { + "action_result": 0, + "get": "ok", + "metrics": { + "a": null, + "c": 0, + "l": [], + "mx": null, + "s": null + }, + "obs_name_given": true, + "types": { + "a": "NoneType", + "c": "int", + "l": "list", + "mx": "NoneType", + "s": "NoneType" + } + }, + "empty_input_no_rows_collected": { + "action_result": [], + "get": "ok", + "metrics": { + "c": 0 + }, + "obs_name_given": true, + "types": { + "c": "int" + } + }, + "err_attribute": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE] Invalid observed metrics. Attribute \"v\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"v\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_attribute_alias_mixed": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE] Invalid observed metrics. Attribute \"v AS raw\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"v AS raw\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_distinct": { + "errorClass": "INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED] Invalid observed metrics. Aggregate expressions with DISTINCT are not allowed in observed metrics, but found: \"count(DISTINCT v) AS d\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"count(DISTINCT v) AS d\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_distinct_sum": { + "errorClass": "INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED] Invalid observed metrics. Aggregate expressions with DISTINCT are not allowed in observed metrics, but found: \"sum(DISTINCT v) AS d\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"sum(DISTINCT v) AS d\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_empty_name_string_api": { + "errorClass": null, + "message": "requirement failed: Name must not be empty", + "messageParameters": {}, + "outcome": "error", + "python_class": "IllegalArgumentException", + "sqlState": null + }, + "err_empty_observation_name": { + "errorClass": "VALUE_NOT_NON_EMPTY_STR", + "message": "[VALUE_NOT_NON_EMPTY_STR] Value for `name` must be a non-empty string, got ''.", + "messageParameters": { + "arg_name": "name", + "arg_value": "" + }, + "outcome": "error", + "python_class": "PySparkValueError", + "sqlState": null + }, + "err_filter": { + "errorClass": "INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_FILTER_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_FILTER_UNSUPPORTED] Invalid observed metrics. Aggregate expression with FILTER predicate are not allowed in observed metrics, but found: \"count(v) FILTER (WHERE (v > 1)) AS f\". SQLSTATE: 42K0E; line 1 pos 0", + "messageParameters": { + "expr": "\"count(v) FILTER (WHERE (v > 1)) AS f\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_first_ignorenulls_ok": { + "get": "ok", + "metrics": { + "f": 2 + }, + "outcome": "ok_unexpected", + "types": { + "f": "int" + } + }, + "err_generator": { + "errorClass": "UNSUPPORTED_GENERATOR.OUTSIDE_SELECT", + "message": "[UNSUPPORTED_GENERATOR.OUTSIDE_SELECT] The generator is not supported: outside the SELECT clause, found: 'CollectMetrics m, [explode(array(1, 2)) AS g#], . SQLSTATE: 42K0E", + "messageParameters": {}, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_get_before_observe": { + "errorClass": "NO_OBSERVE_BEFORE_GET", + "get": "error", + "message": "[NO_OBSERVE_BEFORE_GET] Should observe by calling `DataFrame.observe` before `get`.", + "messageParameters": {}, + "python_class": "PySparkAssertionError", + "sqlState": null + }, + "err_nested_aggregate": { + "errorClass": "INVALID_OBSERVED_METRICS.NESTED_AGGREGATES_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.NESTED_AGGREGATES_UNSUPPORTED] Invalid observed metrics. Nested aggregates are not allowed in observed metrics, but found: \"sum(max(v)) AS n\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"sum(max(v)) AS n\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_no_exprs": { + "errorClass": "CANNOT_BE_EMPTY", + "message": "[CANNOT_BE_EMPTY] At least one exprs must be specified.", + "messageParameters": { + "item": "exprs" + }, + "outcome": "error", + "python_class": "PySparkValueError", + "sqlState": null + }, + "err_no_exprs_string": { + "errorClass": "CANNOT_BE_EMPTY", + "message": "[CANNOT_BE_EMPTY] At least one exprs must be specified.", + "messageParameters": { + "item": "exprs" + }, + "outcome": "error", + "python_class": "PySparkValueError", + "sqlState": null + }, + "err_nondeterministic_plus_agg": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC] Invalid observed metrics. Non-deterministic expression \"(count(1) + rand()) AS r\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"(count(1) + rand()) AS r\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_nondeterministic_rand": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC] Invalid observed metrics. Non-deterministic expression \"rand() AS r\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"rand() AS r\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_observation_name_type": { + "errorClass": "NOT_EXPECTED_TYPE", + "message": "[NOT_EXPECTED_TYPE] Argument `name` should be str, got int.", + "messageParameters": { + "arg_name": "name", + "arg_type": "int", + "expected_type": "str" + }, + "outcome": "error", + "python_class": "PySparkTypeError", + "sqlState": "42K09" + }, + "err_reuse": { + "errorClass": "REUSE_OBSERVATION", + "message": "[REUSE_OBSERVATION] An Observation can be used with a DataFrame only once.", + "messageParameters": {}, + "outcome": "error", + "python_class": "PySparkAssertionError", + "sqlState": null + }, + "err_scalar_subquery": { + "get": "ok", + "metrics": { + "s": 2 + }, + "outcome": "ok_unexpected", + "types": { + "s": "int" + } + }, + "err_unresolved_column": { + "errorClass": "UNRESOLVED_COLUMN.WITH_SUGGESTION", + "message": "[UNRESOLVED_COLUMN.WITH_SUGGESTION] A column, variable, or function parameter with name `nope` cannot be resolved. Did you mean one of the following? [`k`, `v`, `x`]. SQLSTATE: 42703", + "messageParameters": { + "objectName": "`nope`", + "proposal": "`k`, `v`, `x`" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42703" + }, + "err_unresolved_star": { + "errorClass": "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE", + "message": "[INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE] Invalid observed metrics. Attribute \"k\" can only be used as an argument to an aggregate function. SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"k\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "err_window": { + "errorClass": "INVALID_OBSERVED_METRICS.WINDOW_EXPRESSIONS_UNSUPPORTED", + "message": "[INVALID_OBSERVED_METRICS.WINDOW_EXPRESSIONS_UNSUPPORTED] Invalid observed metrics. Window expressions are not allowed in observed metrics, but found: \"sum(v) OVER (ORDER BY k ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS w\". SQLSTATE: 42K0E", + "messageParameters": { + "expr": "\"sum(v) OVER (ORDER BY k ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS w\"" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42K0E" + }, + "exact_under_parallelism": { + "action_result": 100000, + "get": "ok", + "metrics": { + "mx": 99999, + "n": 100000, + "s": 4999950000 + }, + "obs_name_given": true, + "types": { + "mx": "int", + "n": "int", + "s": "int" + } + }, + "explain_does_not_trigger": { + "get": "ok", + "metrics": {}, + "types": {} + }, + "failed_query": { + "action": { + "errorClass": "DIVIDE_BY_ZERO", + "message": "[DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error. SQLSTATE: 22012", + "messageParameters": { + "config": "\"spark.sql.ansi.enabled\"" + }, + "python_class": "ArithmeticException", + "sqlState": "22012" + }, + "get": "ok", + "metrics": "", + "types": "" + }, + "failed_query_then_reuse_df_succeeds": { + "after_fail": { + "get": "ok", + "metrics": "", + "types": "" + }, + "after_success": { + "get": "ok", + "metrics": "", + "types": "" + } + }, + "filter_after_observe": { + "action_result": 1, + "get": "ok", + "metrics": { + "n": 3 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "filter_before_observe": { + "get": "ok", + "metrics": { + "n": 1 + }, + "types": { + "n": "int" + } + }, + "first_then_second_action": { + "after_count": { + "n": 3 + }, + "after_limit1": "" + }, + "flows_no_alias": [ + [ + "x", + { + "count(1)": 3, + "max(v)": 2 + } + ] + ], + "flows_observe_pruned_by_empty_relation": [ + [ + "e", + {} + ] + ], + "flows_string_observe": [ + [ + "nm", + { + "7": 7, + "n": 3, + "s": 3 + } + ] + ], + "flows_two_nested": [ + [ + "a", + { + "n": 3 + } + ], + [ + "b", + { + "s": 3 + } + ] + ], + "get_after_transformation_only_no_action_then_action": { + "after_action": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "before_action": { + "get": "ok", + "metrics": {}, + "types": {} + } + }, + "get_blocks_without_action": { + "get": "ok", + "metrics": {}, + "types": {} + }, + "join_two_observed": { + "a": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "b": { + "get": "ok", + "metrics": { + "s": 3 + }, + "types": { + "s": "int" + } + }, + "count": 3 + }, + "limit_after_observe": { + "action_result": [ + [ + "a", + 1, + "1.5" + ] + ], + "get": "ok", + "metrics": { + "n": 1 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "limit_big_range": { + "action_result": [ + [ + 0 + ], + [ + 1 + ], + [ + 2 + ] + ], + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "obs_name_given": true, + "types": { + "n": "int", + "s": "int" + } + }, + "limit_zero": { + "action_result": [], + "get": "ok", + "metrics": {}, + "obs_name_given": true, + "types": {} + }, + "literal_only": { + "action_result": 3, + "get": "ok", + "metrics": { + "c": 42, + "s": "s" + }, + "obs_name_given": true, + "types": { + "c": "int", + "s": "str" + } + }, + "multiple_actions_first_wins": { + "after_first": { + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "types": { + "n": "int", + "s": "int" + } + }, + "after_more": { + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "types": { + "n": "int", + "s": "int" + } + }, + "counts": [ + 3, + 2, + 3 + ] + }, + "nested_observe_different_names": { + "a": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "b": { + "get": "ok", + "metrics": { + "m": 1 + }, + "types": { + "m": "int" + } + }, + "count": 1 + }, + "nested_observe_same_name_diff_df": { + "errorClass": "DUPLICATED_METRICS_NAME", + "message": "[DUPLICATED_METRICS_NAME] The metric name is not unique: dup. The same name cannot be used for metrics with different results.\nHowever multiple instances of metrics with with same result and name are allowed (e.g. self-joins). SQLSTATE: 42710", + "messageParameters": { + "metricName": "dup" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42710" + }, + "no_alias_names": { + "action_result": 3, + "get": "ok", + "metrics": { + "42": 42, + "count(1)": 3, + "max(v)": 2, + "sum((v + 1))": 5 + }, + "obs_name_given": true, + "types": { + "42": "int", + "count(1)": "int", + "max(v)": "int", + "sum((v + 1))": "int" + } + }, + "nulls_all": { + "action_result": 1, + "get": "ok", + "metrics": { + "c": 0, + "m": null, + "s": null + }, + "obs_name_given": true, + "types": { + "c": "int", + "m": "NoneType", + "s": "NoneType" + } + }, + "observation_in_groupby_downstream": { + "get": "ok", + "metrics": { + "n": 3 + }, + "rows": [ + [ + "a", + 1 + ], + [ + "b", + 1 + ], + [ + "c", + 1 + ] + ], + "types": { + "n": "int" + } + }, + "observation_on_action_collect_schema_only": { + "columns": [ + "k", + "v", + "x" + ], + "get_after_schema": { + "get": "ok", + "metrics": {}, + "types": {} + }, + "schema": "struct" + }, + "observation_through_sort_limit_topn": { + "action_result": [ + [ + 0 + ], + [ + 1 + ], + [ + 2 + ], + [ + 3 + ], + [ + 4 + ] + ], + "get": "ok", + "metrics": { + "n": 20 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "observed_metric_column_name_collision": { + "action_result": 3, + "get": "ok", + "metrics": { + "n": 3 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "ok_aggregate_over_subquery_free_udf_like": { + "get": "ok", + "metrics": { + "s": 3 + }, + "outcome": "ok_unexpected", + "types": { + "s": "int" + } + }, + "ok_approx_count_distinct": { + "get": "ok", + "metrics": { + "a": 2 + }, + "outcome": "ok_unexpected", + "types": { + "a": "int" + } + }, + "ok_count_if": { + "get": "ok", + "metrics": { + "c": 2 + }, + "outcome": "ok_unexpected", + "types": { + "c": "int" + } + }, + "ok_nondeterministic_inside_agg": { + "get": "ok", + "metrics": "", + "outcome": "ok_unexpected", + "types": "" + }, + "python_attrs": { + "has_future": false, + "has_get_async": false, + "has_name": false, + "public": [ + "_error", + "_name", + "_on", + "_result", + "_set_error", + "get" + ] + }, + "same_name_string_same_df_in_both_branches": { + "count": 6 + }, + "self_join_same_df": { + "count": 3, + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "self_union_same_df": { + "count": 6, + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "show_action_limit20": { + "action_result": 0, + "get": "ok", + "metrics": { + "n": 4 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "string_observe_returns_df_only": { + "count": 3, + "type": "DataFrame" + }, + "take_after_observe": { + "action_result": [ + [ + "a", + 1, + "1.5" + ] + ], + "get": "ok", + "metrics": { + "n": 1 + }, + "obs_name_given": true, + "types": { + "n": "int" + } + }, + "types": { + "action_result": 1, + "get": "ok", + "metrics": { + "ai": "1.0", + "l": [ + 1 + ], + "mn": "nan", + "mz": "0.0", + "sd": { + "py_type": "Decimal", + "repr": "1.10" + }, + "si": 1, + "st": [ + 1, + 1 + ] + }, + "obs_name_given": true, + "types": { + "ai": "float", + "l": "list", + "mn": "float", + "mz": "float", + "sd": "Decimal", + "si": "int", + "st": "Row" + } + }, + "union_two_branches_distinct_names": { + "a": { + "get": "ok", + "metrics": { + "n": 3 + }, + "types": { + "n": "int" + } + }, + "b": { + "get": "ok", + "metrics": { + "n": 1 + }, + "types": { + "n": "int" + } + }, + "count": 4 + }, + "union_two_branches_same_name_diff_df": { + "errorClass": "DUPLICATED_METRICS_NAME", + "message": "[DUPLICATED_METRICS_NAME] The metric name is not unique: dup. The same name cannot be used for metrics with different results.\nHowever multiple instances of metrics with with same result and name are allowed (e.g. self-joins). SQLSTATE: 42710", + "messageParameters": { + "metricName": "dup" + }, + "outcome": "error", + "python_class": "AnalysisException", + "sqlState": "42710" + }, + "write_action": { + "get": "ok", + "metrics": { + "n": 3, + "s": 3 + }, + "types": { + "n": "int", + "s": "int" + } + } + }, + "format_version": 1, + "generator": "tests/scripts/gen_observe_goldens.py", + "spark_version": "4.2.0" +} diff --git a/tests/scripts/gen_observe_goldens.py b/tests/scripts/gen_observe_goldens.py new file mode 100644 index 000000000..66d594ffa --- /dev/null +++ b/tests/scripts/gen_observe_goldens.py @@ -0,0 +1,693 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Generate the Dataset.observe / Observation golden fixture (P3-T084). + +Every expectation comes from a REAL Spark 4.2.0 run, never from docs or memory. Each case builds a tiny +DataFrame, attaches an `Observation`, runs an action and records what `Observation.get` returned (values, +Python value types), or the error class / SQLSTATE / message parameters Spark raised. Two client modes are +recorded because they behave differently (see research/t084-observe-semantics.md in krishiv-internal): + + classic : PySpark over the JVM (local[2]); Observation is the JVM `org.apache.spark.sql.Observation`. + connect : PySpark Spark Connect client against a local Connect server (`remote("local[2]")`). + +`get` is called through a 15 s watchdog: `{"get": "blocked"}` means it did not return (it blocks until an +action finishes). Values: ints as ints, doubles as Python `repr` text (`nan`, `inf`), decimals as +`{"py_type": "Decimal", "repr": ...}`. Plan trees and random seeds are stripped from messages; the output +contains no local paths. Cases listed in NONDETERMINISTIC record only the shape, not values. + +Usage (from repo root; needs JDK 21 and `pyspark==4.2.0`): + + ops/scripts/spark-local.sh tests/scripts/gen_observe_goldens.py classic /tmp/observe-classic.json + uv run --no-project --python 3.12 --with 'pyspark[connect]==4.2.0' --with pandas --with pyarrow \\ + python -I tests/scripts/gen_observe_goldens.py connect /tmp/observe-connect.json + python3 tests/scripts/merge_observe_goldens.py /tmp/observe-classic.json /tmp/observe-connect.json \\ + tests/fixtures/spark-4.2.0/observe-goldens.json +""" + +import contextlib +import json +import math +import re +import sys +import threading +from pathlib import Path + +from pyspark.sql import Observation, SparkSession, Window +from pyspark.sql import functions as F + +FORMAT_VERSION = 1 +GENERATOR = "tests/scripts/gen_observe_goldens.py" +NONDETERMINISTIC = {"ok_nondeterministic_inside_agg", "failed_query", "failed_query_then_reuse_df_succeeds"} +MODE = sys.argv[1] # classic | connect +OUT = sys.argv[2] +b = SparkSession.builder.config("spark.ui.enabled", "false").config("spark.sql.shuffle.partitions", "2") +spark = b.remote("local[2]").getOrCreate() if MODE == "connect" else b.master("local[2]").getOrCreate() +print("VERSION", spark.version, file=sys.stderr) + + +def norm(v): + if isinstance(v, float): + if math.isnan(v): + return "nan" + if math.isinf(v): + return "inf" if v > 0 else "-inf" + return repr(v) + if isinstance(v, (list, tuple)): + return [norm(x) for x in v] + if isinstance(v, dict): + return {k: norm(x) for k, x in v.items()} + if v is None: + return None + if isinstance(v, (int, bool, str)): + return v + return {"py_type": type(v).__name__, "repr": str(v)} + + +def err(e): + d = {"python_class": type(e).__name__} + with contextlib.suppress(Exception): + d["errorClass"] = e.getErrorClass() + with contextlib.suppress(Exception): + d["sqlState"] = e.getSqlState() + with contextlib.suppress(Exception): + d["messageParameters"] = {k: str(v) for k, v in (e.getMessageParameters() or {}).items()} + m = str(e) + d["message"] = m.split("\n\n")[0][:600] + return d + + +def with_timeout(f, secs=20): + box = {} + + def run(): + try: + box["v"] = f() + except BaseException as e: + box["e"] = e + + t = threading.Thread(target=run, daemon=True) + t.start() + t.join(secs) + if t.is_alive(): + return ("blocked", None) + if "e" in box: + return ("error", box["e"]) + return ("ok", box["v"]) + + +def getres(o, secs=15): + st, v = with_timeout(lambda: o.get, secs) + if st == "ok": + return {"get": "ok", "metrics": norm(v), "types": {k: type(x).__name__ for k, x in v.items()}} + if st == "blocked": + return {"get": "blocked"} + return {"get": "error", **err(v)} + + +def rows(df): + return [[norm(x) for x in r] for r in df.collect()] + + +spark.range(0).count() +R = {} + + +def case(name): + def deco(f): + try: + R[name] = f() + except BaseException as e: + R[name] = {"outcome": "error", **err(e)} + print("done", name, file=sys.stderr) + + return deco + + +def base(): + return spark.createDataFrame([("a", 1, 1.5), ("b", 2, None), ("c", None, 3.0)], "k string, v int, x double") + + +def empty(): + return base().where("1=0") + + +def run(df_fn, exprs_fn, action=lambda d: d.count(), name=None): + o = Observation(name) if name else Observation() + df = df_fn().observe(o, *exprs_fn()) + res = {"action_result": norm(action(df))} + res.update(getres(o)) + res["obs_name_given"] = bool(name) + return res + + +@case("basic_alias") +def _(): + return run( + base, lambda: [F.count(F.lit(1)).alias("rows"), F.sum("v").alias("sv"), F.max("k").alias("mk")], name="m" + ) + + +@case("no_alias_names") +def _(): + return run(base, lambda: [F.count(F.lit(1)), F.max("v"), F.sum(F.col("v") + 1), F.lit(42)], name="m") + + +@case("literal_only") +def _(): + return run(base, lambda: [F.lit(42).alias("c"), F.lit("s").alias("s")], name="m") + + +@case("composite_expr") +def _(): + return run(base, lambda: [(F.sum("v") + F.avg("x") - F.lit(1)).alias("e"), F.count("*").alias("n")], name="m") + + +@case("empty_input") +def _(): + return run( + empty, + lambda: [ + F.count(F.lit(1)).alias("c"), + F.sum("v").alias("s"), + F.max("k").alias("mx"), + F.avg("x").alias("a"), + F.collect_list("k").alias("l"), + ], + name="m", + ) + + +@case("empty_input_no_rows_collected") +def _(): + return run(empty, lambda: [F.count(F.lit(1)).alias("c")], action=lambda d: d.collect(), name="m") + + +@case("types") +def _(): + return run( + lambda: spark.sql("select cast(1.10 as decimal(5,2)) d, cast('NaN' as double) n, cast(-0.0 as double) z, 1 i"), + lambda: [ + F.sum("d").alias("sd"), + F.max("n").alias("mn"), + F.min("z").alias("mz"), + F.sum("i").alias("si"), + F.avg("i").alias("ai"), + F.collect_list("i").alias("l"), + F.struct(F.max("i"), F.min("i")).alias("st"), + ], + name="m", + ) + + +@case("nulls_all") +def _(): + return run( + lambda: spark.sql("select cast(null as int) v"), + lambda: [F.sum("v").alias("s"), F.count("v").alias("c"), F.max("v").alias("m")], + name="m", + ) + + +@case("anonymous_observation") +def _(): + o = Observation() + df = base().observe(o, F.count(F.lit(1)).alias("n")) + df.count() + return {"name_is_str": isinstance(o._name, str), "name_len": len(o._name) if o._name else None, **getres(o)} + + +# errors +def err_case(exprs_fn, df_fn=base): + def f(): + o = Observation("m") + df = df_fn().observe(o, *exprs_fn()) + df.count() # analysis error surfaces at action (lazily) + return {"outcome": "ok_unexpected", **getres(o)} + + return f + + +for nm, fn in { + "err_attribute": lambda: [F.col("v")], + "err_attribute_alias_mixed": lambda: [F.count(F.lit(1)).alias("c"), F.col("v").alias("raw")], + "err_nondeterministic_rand": lambda: [F.rand().alias("r")], + "ok_nondeterministic_inside_agg": lambda: [F.sum(F.rand()).alias("r")], + "err_nondeterministic_plus_agg": lambda: [(F.count(F.lit(1)) + F.rand()).alias("r")], + "err_nested_aggregate": lambda: [F.sum(F.max("v")).alias("n")], + "err_distinct": lambda: [F.countDistinct("v").alias("d")], + "err_distinct_sum": lambda: [F.sum_distinct(F.col("v")).alias("d")], + "err_filter": lambda: [F.expr("count(v) filter (where v > 1)").alias("f")], + "err_window": lambda: [F.sum("v").over(Window.orderBy("k")).alias("w")], + "err_generator": lambda: [F.explode(F.array(F.lit(1), F.lit(2))).alias("g")], + "err_unresolved_column": lambda: [F.sum("nope").alias("s")], + "err_scalar_subquery": lambda: [F.expr("(select max(id) from range(3))").alias("s")], + "ok_aggregate_over_subquery_free_udf_like": lambda: [F.sum(F.length("k")).alias("s")], + "err_unresolved_star": lambda: [F.expr("count(*)").alias("c"), F.col("*")], + "ok_approx_count_distinct": lambda: [F.approx_count_distinct("v").alias("a")], + "err_first_ignorenulls_ok": lambda: [F.first("v", True).alias("f")], + "ok_count_if": lambda: [F.expr("count_if(v>0)").alias("c")], +}.items(): + case(nm)(err_case(fn)) + + +@case("err_empty_name_string_api") +def _(): + df = base().observe("", F.count(F.lit(1)).alias("n")) + df.count() + return {"outcome": "ok_unexpected"} + + +@case("err_empty_observation_name") +def _(): + Observation("") + return {"outcome": "ok_unexpected"} + + +@case("err_observation_name_type") +def _(): + Observation(5) + return {"outcome": "ok_unexpected"} + + +@case("err_no_exprs") +def _(): + Observation("m") + base().observe(Observation("m")) + return {"outcome": "ok_unexpected"} + + +@case("err_no_exprs_string") +def _(): + base().observe("m") + return {"outcome": "ok_unexpected"} + + +@case("err_get_before_observe") +def _(): + o = Observation("m") + return getres(o, 3) + + +@case("err_reuse") +def _(): + o = Observation("m") + base().observe(o, F.count(F.lit(1)).alias("n")) + base().observe(o, F.count(F.lit(1)).alias("n")) + return {"outcome": "ok_unexpected"} + + +@case("get_blocks_without_action") +def _(): + o = Observation("m") + base().observe(o, F.count(F.lit(1)).alias("n")) + return getres(o, 4) + + +@case("get_after_transformation_only_no_action_then_action") +def _(): + o = Observation("m") + df = base().observe(o, F.count(F.lit(1)).alias("n")) + r1 = getres(o, 3) + df.count() + r2 = getres(o, 10) + return {"before_action": r1, "after_action": r2} + + +@case("multiple_actions_first_wins") +def _(): + o = Observation("m") + df = base().observe(o, F.count(F.lit(1)).alias("n"), F.sum("v").alias("s")) + c1 = df.count() + r1 = getres(o) + c2 = df.where("v is not null").count() + c3 = df.collect() + r2 = getres(o) + return {"counts": [c1, c2, len(c3)], "after_first": r1, "after_more": r2} + + +@case("count_action_prunes_columns") +def _(): + # df.count() prunes columns; metric on v still computed? (CollectMetrics keeps referenced cols) + return run(base, lambda: [F.sum("v").alias("s")], action=lambda d: d.count(), name="m") + + +@case("filter_after_observe") +def _(): + return run(base, lambda: [F.count(F.lit(1)).alias("n")], action=lambda d: d.where("v = 1").count(), name="m") + + +@case("filter_before_observe") +def _(): + o = Observation("m") + df = base().where("v = 1").observe(o, F.count(F.lit(1)).alias("n")) + df.collect() + return getres(o) + + +@case("limit_after_observe") +def _(): + return run(base, lambda: [F.count(F.lit(1)).alias("n")], action=lambda d: d.limit(1).collect(), name="m") + + +@case("take_after_observe") +def _(): + return run(base, lambda: [F.count(F.lit(1)).alias("n")], action=lambda d: d.take(1), name="m") + + +@case("limit_zero") +def _(): + return run(base, lambda: [F.count(F.lit(1)).alias("n")], action=lambda d: d.limit(0).collect(), name="m") + + +@case("limit_big_range") +def _(): + return run( + lambda: spark.range(0, 100000, 1, 4).toDF("v"), + lambda: [F.count(F.lit(1)).alias("n"), F.sum("v").alias("s")], + action=lambda d: d.limit(3).collect(), + name="m", + ) + + +@case("exact_under_parallelism") +def _(): + return run( + lambda: spark.range(0, 100000, 1, 8).toDF("v"), + lambda: [F.count(F.lit(1)).alias("n"), F.sum("v").alias("s"), F.max("v").alias("mx")], + name="m", + ) + + +@case("cache") +def _(): + o = Observation("m") + df = base().observe(o, F.count(F.lit(1)).alias("n")).cache() + a = df.count() + r1 = getres(o) + b2 = df.count() + r2 = getres(o) + return {"counts": [a, b2], "after_first": r1, "after_second": r2} + + +@case("cache_then_new_observe_on_cached") +def _(): + df = base().cache() + df.count() + o = Observation("m") + d2 = df.observe(o, F.count(F.lit(1)).alias("n")) + d2.collect() + return getres(o) + + +@case("union_two_branches_same_name_diff_df") +def _(): + a = base().observe("dup", F.count(F.lit(1)).alias("n")) + b2 = base().observe("dup", F.count(F.lit(1)).alias("n")) + a.union(b2).count() + return {"outcome": "ok_unexpected"} + + +@case("union_two_branches_distinct_names") +def _(): + oa, ob = Observation("a"), Observation("b") + a = base().observe(oa, F.count(F.lit(1)).alias("n")) + b2 = base().where("v=1").observe(ob, F.count(F.lit(1)).alias("n")) + c = a.union(b2).count() + return {"count": c, "a": getres(oa), "b": getres(ob)} + + +@case("self_join_same_df") +def _(): + o = Observation("m") + d = base().observe(o, F.count(F.lit(1)).alias("n")) + c = d.alias("l").join(d.alias("r"), F.col("l.k") == F.col("r.k")).count() + return {"count": c, **getres(o)} + + +@case("self_union_same_df") +def _(): + o = Observation("m") + d = base().observe(o, F.count(F.lit(1)).alias("n")) + c = d.union(d).count() + return {"count": c, **getres(o)} + + +@case("join_two_observed") +def _(): + oa, ob = Observation("a"), Observation("b") + a = base().observe(oa, F.count(F.lit(1)).alias("n")) + b2 = base().observe(ob, F.sum("v").alias("s")) + c = a.join(b2, "k").count() + return {"count": c, "a": getres(oa), "b": getres(ob)} + + +@case("same_name_string_same_df_in_both_branches") +def _(): + d = base().observe("dup", F.count(F.lit(1)).alias("n")) + return {"count": d.union(d).count()} + + +@case("nested_observe_different_names") +def _(): + oa, ob = Observation("a"), Observation("b") + d = base().observe(oa, F.count(F.lit(1)).alias("n")).where("v=1").observe(ob, F.count(F.lit(1)).alias("m")) + c = d.count() + return {"count": c, "a": getres(oa), "b": getres(ob)} + + +@case("nested_observe_same_name_diff_df") +def _(): + d = base().observe("dup", F.count(F.lit(1)).alias("n")).where("v=1").observe("dup", F.count(F.lit(1)).alias("n")) + return {"count": d.count()} + + +@case("observed_metric_column_name_collision") +def _(): + return run(base, lambda: [F.count(F.lit(1)).alias("n"), F.sum("v").alias("n")], name="m") + + +@case("failed_query") +def _(): + spark.conf.set("spark.sql.ansi.enabled", "true") + try: + o = Observation("m") + d = spark.range(5).toDF("v").observe(o, F.count(F.lit(1)).alias("n")).selectExpr("1/(v-2) as r") + try: + d.collect() + a = "ok_unexpected" + except Exception as e: + a = err(e) + return {"action": a, **getres(o, 15)} + finally: + spark.conf.set("spark.sql.ansi.enabled", "true") + + +@case("failed_query_then_reuse_df_succeeds") +def _(): + o = Observation("m") + base0 = spark.range(5).toDF("v").observe(o, F.count(F.lit(1)).alias("n")) + with contextlib.suppress(Exception): + base0.selectExpr("1/(v-2) as r").collect() + r1 = getres(o, 8) + base0.count() + return {"after_fail": r1, "after_success": getres(o, 8)} + + +@case("observation_on_action_collect_schema_only") +def _(): + o = Observation("m") + d = base().observe(o, F.count(F.lit(1)).alias("n")) + s = d.schema.simpleString() + c = d.columns + return {"schema": s, "columns": c, "get_after_schema": getres(o, 3)} + + +@case("explain_does_not_trigger") +def _(): + o = Observation("m") + d = base().observe(o, F.count(F.lit(1)).alias("n")) + d.explain() + return getres(o, 3) + + +@case("write_action") +def _(): + import tempfile + + o = Observation("m") + d = base().observe(o, F.count(F.lit(1)).alias("n"), F.sum("v").alias("s")) + p = tempfile.mkdtemp() + "/o" + d.write.parquet(p) + return getres(o) + + +@case("string_observe_returns_df_only") +def _(): + d = base().observe("nm", F.count(F.lit(1)).alias("n")) + return {"count": d.count(), "type": type(d).__name__} + + +@case("observation_in_groupby_downstream") +def _(): + o = Observation("m") + d = base().observe(o, F.count(F.lit(1)).alias("n")).groupBy("k").count() + return {"rows": sorted(rows(d)), **getres(o)} + + +@case("observation_through_sort_limit_topn") +def _(): + return run( + lambda: spark.range(0, 1000, 1, 4).toDF("v"), + lambda: [F.count(F.lit(1)).alias("n")], + action=lambda d: d.orderBy("v").limit(5).collect(), + name="m", + ) + + +@case("show_action_limit20") +def _(): + return run( + lambda: spark.range(0, 1000, 1, 4).toDF("v"), + lambda: [F.count(F.lit(1)).alias("n")], + action=lambda d: d.show(3) or 0, + name="m", + ) + + +@case("python_attrs") +def _(): + o = Observation("m") + return { + "has_get_async": hasattr(o, "get_async"), + "has_future": hasattr(o, "future"), + "has_name": hasattr(o, "name"), + "public": [a for a in dir(o) if not a.startswith("__")], + } + + +@case("aggregate_with_alias_inside_nested_alias") +def _(): + return run(base, lambda: [F.sum("v").alias("a").alias("b")], name="m") + + +@case("ansi_off_overflow_sum") +def _(): + spark.conf.set("spark.sql.ansi.enabled", "false") + try: + return run( + lambda: spark.sql("select cast(9223372036854775807 as bigint) v union all select 1"), + lambda: [F.sum("v").alias("s"), (F.max("v") + 1).alias("p1")], + name="m", + ) + finally: + spark.conf.set("spark.sql.ansi.enabled", "true") + + +@case("ansi_on_overflow_sum") +def _(): + spark.conf.set("spark.sql.ansi.enabled", "true") + return run( + lambda: spark.sql("select cast(9223372036854775807 as bigint) v union all select 1"), + lambda: [F.sum("v").alias("s")], + name="m", + ) + + +@case("ansi_on_cast_error_inside_metric") +def _(): + return run(base, lambda: [F.sum(F.col("k").cast("int")).alias("s")], name="m") + + +@case("conditional_metric_no_raise_unselected") +def _(): + return run( + lambda: spark.sql("select * from values (0),(2) as t(v)"), + lambda: [F.sum(F.when(F.col("v") != 0, F.lit(10) / F.col("v"))).alias("s")], + name="m", + ) + + +# ---- second-action probes (first-wins in classic, last-wins in Connect) ---- +@case("first_then_second_action") +def _(): + o = Observation("m") + df = base().observe(o, F.count(F.lit(1)).alias("n")) + df.limit(1).collect() + first = dict(o.get) + df.count() + second = dict(o.get) + return {"after_limit1": first, "after_count": second} + + +if MODE == "connect": + + @case("flows_string_observe") + def _(): + d = base().observe("nm", F.count(F.lit(1)).alias("n"), F.sum("v").alias("s"), F.lit(7)) + d.toPandas() + return d.executionInfo.flows + + @case("flows_two_nested") + def _(): + d = base().observe("a", F.count(F.lit(1)).alias("n")).observe("b", F.sum("v").alias("s")) + d.toPandas() + return d.executionInfo.flows + + @case("flows_no_alias") + def _(): + d = base().observe("x", F.count(F.lit(1)), F.max("v")) + d.toPandas() + return d.executionInfo.flows + + @case("flows_observe_pruned_by_empty_relation") + def _(): + d = base().observe("e", F.count(F.lit(1)).alias("n")).where("1=0") + d.toPandas() + return d.executionInfo.flows + + +def _scrub(m): + m = re.sub(r"rand\(-?\d+\)", "rand()", m) + m = re.sub(r"#\d+L?", "#", m) + return re.sub(r"\], \d+", "], ", m) + + +def _clean_msg(m): + return _scrub(m.split("\n==")[0].split(";\n")[0].strip()) + + +def clean(name, v): + if isinstance(v, dict): + out = {} + for k, x in v.items(): + if k == "message" and isinstance(x, str): + out[k] = _clean_msg( + x if "AssertionError" not in x else "java.lang.AssertionError in PythonSQLUtils.toPyRow" + ) + elif k == "messageParameters" and isinstance(x, dict): + out[k] = {a: _scrub(b) for a, b in x.items() if a != "plan"} + elif name == "first_then_second_action" and k == "after_limit1" and MODE == "connect": + out[k] = "" # Connect: 1 in four isolated sessions, 3 once inside the full suite + elif k in ("metrics", "types") and name in NONDETERMINISTIC: + out[k] = "" + else: + out[k] = clean(name, x) + return out + return v + + +R = {k: clean(k, v) for k, v in R.items()} +Path(OUT).write_text( + json.dumps( + { + "format_version": FORMAT_VERSION, + "generator": GENERATOR, + "mode": MODE, + "spark_version": spark.version, + "cases": R, + }, + indent=1, + sort_keys=True, + ) +) +spark.stop() diff --git a/tests/scripts/merge_observe_goldens.py b/tests/scripts/merge_observe_goldens.py new file mode 100644 index 000000000..3d534c1df --- /dev/null +++ b/tests/scripts/merge_observe_goldens.py @@ -0,0 +1,28 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Merge the classic and Connect runs of gen_observe_goldens.py into one fixture (P3-T084). + +Usage: python3 tests/scripts/merge_observe_goldens.py CLASSIC.json CONNECT.json OUT.json +""" + +import json +import sys +from pathlib import Path + + +def main() -> None: + classic = json.loads(Path(sys.argv[1]).read_text()) + connect = json.loads(Path(sys.argv[2]).read_text()) + assert classic["spark_version"] == connect["spark_version"] + merged = { + "format_version": classic["format_version"], + "generator": classic["generator"], + "spark_version": classic["spark_version"], + "classic": classic["cases"], + "connect": connect["cases"], + } + Path(sys.argv[3]).write_text(json.dumps(merged, indent=1, sort_keys=True) + "\n") + + +if __name__ == "__main__": + main() From 8784fbd615609e958e9a716fa1705d960701ea39 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 00:00:38 +0000 Subject: [PATCH 02/15] plan: PlanNode::CollectMetrics and its analysis (P3-T084 unit 2) Dataset.observe as a logical node: rows pass through, the metrics resolve as a global aggregate over the input, and Spark 4.2.0's CheckAnalysis classes apply (INVALID_OBSERVED_METRICS.*, DUPLICATED_METRICS_NAME across different relations, UNSUPPORTED_GENERATOR.OUTSIDE_SELECT). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../src/analyzer/exprs/aggregates.rs | 18 +- crates/krishiv-plan/src/analyzer/exprs/mod.rs | 8 +- .../krishiv-plan/src/analyzer/exprs/node.rs | 2 +- crates/krishiv-plan/src/analyzer/mod.rs | 1 + crates/krishiv-plan/src/analyzer/observe.rs | 171 +++++++++++++++ crates/krishiv-plan/src/analyzer/resolve.rs | 13 ++ crates/krishiv-plan/src/cache.rs | 10 + crates/krishiv-plan/src/display.rs | 8 + crates/krishiv-plan/src/encode.rs | 5 + crates/krishiv-plan/src/lib.rs | 1 + crates/krishiv-plan/src/plan.rs | 27 ++- crates/krishiv-plan/src/tests/analyzer.rs | 4 +- crates/krishiv-plan/src/tests/inventory.rs | 9 + crates/krishiv-plan/src/tests/mod.rs | 1 + crates/krishiv-plan/src/tests/observe.rs | 199 ++++++++++++++++++ crates/krishiv-plan/src/tests/props.rs | 12 +- crates/krishiv-plan/src/walk.rs | 6 + 17 files changed, 484 insertions(+), 11 deletions(-) create mode 100644 crates/krishiv-plan/src/analyzer/observe.rs create mode 100644 crates/krishiv-plan/src/tests/observe.rs diff --git a/crates/krishiv-plan/src/analyzer/exprs/aggregates.rs b/crates/krishiv-plan/src/analyzer/exprs/aggregates.rs index 2e2619842..c82b3254f 100644 --- a/crates/krishiv-plan/src/analyzer/exprs/aggregates.rs +++ b/crates/krishiv-plan/src/analyzer/exprs/aggregates.rs @@ -44,7 +44,13 @@ impl Resolver<'_> { let mut checked: Vec = kids[..filter_at.min(kids.len())].to_vec(); checked.extend(orders.iter().cloned()); if !window_fn { - call_checks(entry, e, &checked, filter_at.min(checked.len()))?; + call_checks( + entry, + e, + &checked, + filter_at.min(checked.len()), + self.observing, + )?; } let wrong_arity = || { AnalysisError::new( @@ -285,7 +291,13 @@ impl Resolver<'_> { let mut checked: Vec = kids[..filter_at.min(kids.len())].to_vec(); checked.extend(orders.iter().cloned()); if !window_fn { - call_checks(entry, e, &checked, filter_at.min(checked.len()))?; + call_checks( + entry, + e, + &checked, + filter_at.min(checked.len()), + self.observing, + )?; } // The value and the delimiter, implicitly cast (Spark's arity // check, then `listagg`'s own delimiter checks in the registry). @@ -617,7 +629,7 @@ impl Resolver<'_> { } // The nesting and window checks see the filter and the ordering key // too (a filter sits between the arguments and the ordering keys). - call_checks(entry, e, kids, kids.len())?; + call_checks(entry, e, kids, kids.len(), self.observing)?; let filter = kids.get(n).filter(|_| f.filter.is_some()); let wstart = kids.len() - f.within_group.len(); // The lowered call's SQL as Spark quotes it in errors: `name(args)` diff --git a/crates/krishiv-plan/src/analyzer/exprs/mod.rs b/crates/krishiv-plan/src/analyzer/exprs/mod.rs index 87100bc2f..ae6120641 100644 --- a/crates/krishiv-plan/src/analyzer/exprs/mod.rs +++ b/crates/krishiv-plan/src/analyzer/exprs/mod.rs @@ -762,7 +762,8 @@ impl Resolver<'_> { /// of an aggregate calls a nondeterministic function. A call of a SQL /// function is as deterministic as its body, whatever its arguments. fn nondeterministic_arguments(&mut self, e: &Arc, args: &[Bound]) { - if self.pending_nondeterministic.is_some() { + // `sum(rand())` is a legal observed metric. + if self.observing || self.pending_nondeterministic.is_some() { return; } let Some(bad) = args.iter().find(|k| { @@ -860,6 +861,7 @@ fn call_checks( e: &Arc, kids: &[Bound], n: usize, + observing: bool, ) -> Result<(), AnalysisError> { if entry.kind == FunctionKind::Window { let shown = with_children(e, kids.iter().map(Bound::shown).collect()); @@ -870,7 +872,9 @@ fn call_checks( )); } if entry.kind == FunctionKind::Aggregate { - if let Some(inner) = kids.iter().find_map(|k| k.agg) { + // Observed metrics report a nested aggregate as their own class + // (`observe` checks them). + if !observing && let Some(inner) = kids.iter().find_map(|k| k.agg) { return Err(AnalysisError::new(NestedAggregateFunction {}, inner)); } if kids[..n].iter().any(|k| k.win.is_some()) { diff --git a/crates/krishiv-plan/src/analyzer/exprs/node.rs b/crates/krishiv-plan/src/analyzer/exprs/node.rs index 0f3e46ce1..57acfa636 100644 --- a/crates/krishiv-plan/src/analyzer/exprs/node.rs +++ b/crates/krishiv-plan/src/analyzer/exprs/node.rs @@ -619,7 +619,7 @@ impl Resolver<'_> { return self.quantile_call(&entry, e, f, &kids, n, window_fn); } if !window_fn { - call_checks(&entry, e, &kids, n)?; + call_checks(&entry, e, &kids, n, self.observing)?; } // `struct(a, ...)`: children with names keep them (a column // contributes its name, a struct field its field name), so a diff --git a/crates/krishiv-plan/src/analyzer/mod.rs b/crates/krishiv-plan/src/analyzer/mod.rs index c59908c28..e273b51fc 100644 --- a/crates/krishiv-plan/src/analyzer/mod.rs +++ b/crates/krishiv-plan/src/analyzer/mod.rs @@ -98,6 +98,7 @@ mod hof; mod instance; mod intervals; pub(crate) mod legacy; +mod observe; mod ops; mod pivot; pub(crate) mod prefix; diff --git a/crates/krishiv-plan/src/analyzer/observe.rs b/crates/krishiv-plan/src/analyzer/observe.rs new file mode 100644 index 000000000..7ab3e641b --- /dev/null +++ b/crates/krishiv-plan/src/analyzer/observe.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `Dataset.observe(name, metrics)` (Spark's `CollectMetrics`): the +//! metrics are a global aggregate over the input that the rows pass +//! through, so they resolve as the select list of an aggregate with no +//! grouping, and `CheckAnalysis.checkMetric` (4.2.0) limits what they may +//! hold. Rows are untouched: the node's output is its input's. +//! +//! Order of errors, as Spark raises them: the analyzer's (an unresolved +//! column, a generator) first, then the missing name, then per metric in +//! order the first violation of a pre-order walk (window, nested +//! aggregate, DISTINCT, FILTER, a column or a non-deterministic call +//! outside any aggregate). The expression shown is the whole aliased +//! metric. A name used by two relations that are not the same relation +//! (equal field by field, `plan_id` included) is `DUPLICATED_METRICS_NAME`; +//! the same relation reached twice (a self-join, a self-union) is legal. + +use std::sync::Arc; + +use krishiv_common::error::{ + DuplicatedMetricsName, InvalidObservedMetricsAggregateExpressionWithDistinctUnsupported, + InvalidObservedMetricsAggregateExpressionWithFilterUnsupported, + InvalidObservedMetricsMissingName, InvalidObservedMetricsNestedAggregatesUnsupported, + InvalidObservedMetricsNonAggregateFuncArgIsAttribute, + InvalidObservedMetricsNonAggregateFuncArgIsNonDeterministic, + InvalidObservedMetricsWindowExpressionsUnsupported, UnsupportedGeneratorOutsideSelect, +}; + +use super::AnalysisError; +use super::bind::addr; +use super::exprs::pretty; +use super::ops::has_generator; +use super::resolve::{Env, Rel, Resolver}; +use crate::{Expr, ExprNode, FunctionKind, LogicalPlan, PlanNode, ResolvedExpr}; + +/// What an observed metric may not hold, in Spark's order of checks. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum Violation { + Window, + Nested, + Distinct, + Filter, + Attribute, + NonDeterministic, +} + +/// The first violation in a pre-order walk of `metric`, `None` when the +/// metric is allowed. +fn violation(metric: &Arc) -> Option { + let mut stack = vec![(Arc::clone(metric), false)]; + while let Some((e, seen)) = stack.pop() { + let mut inside = seen; + match &e.node { + ExprNode::Window(_) => return Some(Violation::Window), + ExprNode::Resolved(ResolvedExpr::Function(f)) if f.kind == FunctionKind::Aggregate => { + if seen { + return Some(Violation::Nested); + } + if f.distinct { + return Some(Violation::Distinct); + } + if f.filter.is_some() { + return Some(Violation::Filter); + } + inside = true; + } + ExprNode::Resolved(ResolvedExpr::Column(_)) if !seen => { + return Some(Violation::Attribute); + } + ExprNode::Resolved(ResolvedExpr::Alias { .. }) => {} + _ if !seen && !crate::raises::is_deterministic(&e) => { + return Some(Violation::NonDeterministic); + } + _ => {} + } + let mut kids = Vec::new(); + e.for_each_child(&mut |k| kids.push(Arc::clone(k))); + stack.extend(kids.into_iter().rev().map(|k| (k, inside))); + } + None +} + +impl Resolver<'_> { + /// Resolves `metrics` over `input` as a global aggregate would and + /// keeps the node with them bound; its output is the input's. + pub(super) fn collect_metrics( + &mut self, + p: &Arc, + input: &Arc, + name: &str, + metrics: &[Arc], + env: &Arc, + ) -> Result { + let o = p.origin; + let c = self.child(input, env)?; + self.observing = true; + let items = self.select_items(&c, metrics, env, "CollectMetrics"); + self.observing = false; + let items = items?; + let zone = self.ctx.conf.session_zone; + let shown = |b: &super::exprs::Bound| format!("\"{}\"", pretty(&b.shown(), zone)); + if let Some(b) = items.iter().find(|b| has_generator(&b.expr)) { + return Err(AnalysisError::new( + UnsupportedGeneratorOutsideSelect { + plan: format!("CollectMetrics {name}, [{}]", pretty(&b.shown(), zone)), + }, + o, + )); + } + if name.is_empty() || items.is_empty() { + let listed: Vec = items.iter().map(|b| pretty(&b.shown(), zone)).collect(); + return Err(AnalysisError::new( + InvalidObservedMetricsMissingName { + operator: format!("CollectMetrics {name}, [{}]", listed.join(", ")), + }, + o, + )); + } + for b in &items { + let Some(v) = violation(&b.expr) else { + continue; + }; + let expr = shown(b); + let e: krishiv_common::error::KrishivError = match v { + Violation::Window => { + InvalidObservedMetricsWindowExpressionsUnsupported { expr }.into() + } + Violation::Nested => { + InvalidObservedMetricsNestedAggregatesUnsupported { expr }.into() + } + Violation::Distinct => { + InvalidObservedMetricsAggregateExpressionWithDistinctUnsupported { expr }.into() + } + Violation::Filter => { + InvalidObservedMetricsAggregateExpressionWithFilterUnsupported { expr }.into() + } + Violation::Attribute => { + InvalidObservedMetricsNonAggregateFuncArgIsAttribute { expr }.into() + } + Violation::NonDeterministic => { + InvalidObservedMetricsNonAggregateFuncArgIsNonDeterministic { expr }.into() + } + }; + return Err(AnalysisError::new(e, o)); + } + // One relation per name: the same relation (its canonical + // representative) may appear again, another may not. + let key = addr(&*self.rep(p)); + match self.observed.get(name) { + Some(first) if *first != key => { + return Err(AnalysisError::new( + DuplicatedMetricsName { + metric_name: name.to_owned(), + }, + o, + )); + } + Some(_) => {} + None => { + self.observed.insert(name.to_owned(), key); + } + } + let node = PlanNode::CollectMetrics { + input: Arc::clone(&c.plan), + name: name.to_owned(), + metrics: items.into_iter().map(|b| b.expr).collect(), + }; + let mut r = self.make(o, node, c.out.to_vec(), &[&c]); + r.hidden = Arc::clone(&c.hidden); + Ok(r) + } +} diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index 208873532..0aeeb8524 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -190,6 +190,8 @@ pub(super) struct Resolver<'a> { /// (`encode::identity_keys`): resolution does not depend on how the /// input shares subtrees, so equal input keys give equal plans. pub(super) canon: HashMap>, + /// The relation (canonical address) that first used each observation name (`observe`). + pub(super) observed: HashMap, /// The identity key of every relation of the input, kept from the first `canonicalize` (the root's). pub(super) keys: HashMap, /// The input holds a node the prefix cache does not serve (a resolved scan, SQL text, a read by path). @@ -232,6 +234,8 @@ pub(super) struct Resolver<'a> { /// The first aggregate over a nondeterministic argument: Spark checks /// it after the operators' other rules, so any other error wins. pub(super) pending_nondeterministic: Option, + /// Binding the metrics of an `observe` (Spark checks them in `CheckAnalysis`, not as an aggregate). + pub(super) observing: bool, /// The attributes that hold a `CHAR(n)` / `VARCHAR(n)` column, by id, with the catalog string they were /// declared with (see `char_varchar`). pub(super) declared: HashMap>, @@ -269,6 +273,8 @@ pub(super) fn resolve( memo: HashMap::new(), closed: HashMap::new(), canon: HashMap::new(), + observed: HashMap::new(), + observing: false, keys: HashMap::new(), uncacheable: false, warm_started: false, @@ -665,6 +671,7 @@ impl Resolver<'_> { | PlanNode::Limit { .. } | PlanNode::Offset { .. } | PlanNode::Repartition { .. } + | PlanNode::CollectMetrics { .. } ); let plan = Arc::new(LogicalPlan::new(node).with_origin(origin)); let mut r = self.wrap(plan, origin.plan_id, out, kids); @@ -940,6 +947,11 @@ impl Resolver<'_> { } => self.hint(p, input, name, params, env), PlanNode::PipeOperator { input } => self.pipe_operator(input, env), PlanNode::Repartition { input, exprs } => self.repartition(p, input, exprs, env), + PlanNode::CollectMetrics { + input, + name, + metrics, + } => self.collect_metrics(p, input, name, metrics, env), PlanNode::TableFunction { name, args } => self.table_function(p, name, args, env), // Listed, not wildcarded: a new `PlanNode` variant must be // placed here or in `leaf_node` to compile. @@ -1120,6 +1132,7 @@ impl Resolver<'_> { | PlanNode::Hint { .. } | PlanNode::PipeOperator { .. } | PlanNode::Repartition { .. } + | PlanNode::CollectMetrics { .. } | PlanNode::Union { .. } | PlanNode::Distinct { .. } | PlanNode::Pivot { .. } diff --git a/crates/krishiv-plan/src/cache.rs b/crates/krishiv-plan/src/cache.rs index 9dc31b368..a6a579ef7 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -821,6 +821,16 @@ impl Encoder { self.exprs(exprs, bias); self.plan(input, bias); } + PlanNode::CollectMetrics { + input, + name, + metrics, + } => { + self.tag("CollectMetrics"); + self.text(name); + self.exprs(metrics, bias); + self.plan(input, bias); + } PlanNode::JoinHint { input, strategy } => { self.tag("JoinHint"); self.debug(strategy); diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index be51c22d1..4cd624cd7 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -356,6 +356,14 @@ impl Display for NodeLine<'_> { } PlanNode::PipeOperator { .. } => f.write_str("PipeOperator"), PlanNode::Repartition { exprs, .. } => write!(f, "Repartition ({})", List(exprs)), + PlanNode::CollectMetrics { name, metrics, .. } => { + write!( + f, + "CollectMetrics {} ({})", + Name(std::slice::from_ref(name)), + List(metrics) + ) + } PlanNode::JoinHint { strategy, .. } => { write!(f, "JoinHint (strategy={})", strategy.name()) } diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 582c68f8f..fb5485505 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -113,6 +113,7 @@ pub(crate) fn plan_tag(node: &PlanNode) -> u8 { PlanNode::Generate { .. } => 33, PlanNode::PipeOperator { .. } => 34, PlanNode::Repartition { .. } => 35, + PlanNode::CollectMetrics { .. } => 36, } } @@ -900,6 +901,10 @@ impl<'k, 'a> Encoder<'k, 'a> { | PlanNode::PipeOperator { .. } | PlanNode::Distinct { .. } => {} PlanNode::Repartition { exprs, .. } => self.exprs(exprs), + PlanNode::CollectMetrics { name, metrics, .. } => { + self.str(name); + self.exprs(metrics); + } PlanNode::JoinHint { strategy, .. } => { self.str(strategy.name()); } diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 16de1d556..7d44552e0 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -80,6 +80,7 @@ //! | `PlanNode::Hint` | U | - | - | //! | `PlanNode::PipeOperator` | U | - | - | //! | `PlanNode::Repartition` | U, R | - | - | +//! | `PlanNode::CollectMetrics` | U, R | - | `observe` | //! | `PlanNode::JoinHint` | R | - | - | //! | `PlanNode::Union` | U, R | - | - | //! | `PlanNode::Distinct` | U | - | - | diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 4a377a245..64a8c99be 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -239,6 +239,18 @@ pub enum PlanNode { /// The partitioning expressions. exprs: Vec>, }, + /// `Dataset.observe(name, metrics)` (Spark's `CollectMetrics`): rows + /// pass through unchanged while the metrics (aggregate expressions, as + /// a global aggregate with no grouping) are folded over the rows that + /// flow through; the result is reported under `name`, not as rows. + CollectMetrics { + /// Input. + input: Arc, + /// The observation name. + name: String, + /// The metric expressions (aliased once resolved). + metrics: Vec>, + }, /// A join strategy hint on a relation (Spark's `ResolvedHint`): a join /// this relation is a side of prefers `strategy`. Rows pass through. JoinHint { @@ -981,6 +993,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::PipeOperator { input } | PlanNode::Repartition { input, .. } + | PlanNode::CollectMetrics { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } | PlanNode::Expand { input, .. } @@ -1031,6 +1044,7 @@ impl LogicalPlan { | PlanNode::Hint { input, .. } | PlanNode::PipeOperator { input } | PlanNode::Repartition { input, .. } + | PlanNode::CollectMetrics { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Distinct { input } | PlanNode::Expand { input, .. } @@ -1070,7 +1084,11 @@ impl LogicalPlan { | PlanNode::Union { .. } | PlanNode::SetOp { .. } | PlanNode::Distinct { .. } => {} - PlanNode::Hint { params, .. } | PlanNode::Repartition { exprs: params, .. } => { + PlanNode::Hint { params, .. } + | PlanNode::Repartition { exprs: params, .. } + | PlanNode::CollectMetrics { + metrics: params, .. + } => { params.iter().for_each(f); } PlanNode::Sql { args, .. } => match args { @@ -1154,7 +1172,11 @@ impl LogicalPlan { | PlanNode::Union { .. } | PlanNode::SetOp { .. } | PlanNode::Distinct { .. } => {} - PlanNode::Hint { params, .. } | PlanNode::Repartition { exprs: params, .. } => { + PlanNode::Hint { params, .. } + | PlanNode::Repartition { exprs: params, .. } + | PlanNode::CollectMetrics { + metrics: params, .. + } => { params.iter_mut().for_each(f); } PlanNode::Sql { args, .. } => match args { @@ -1253,6 +1275,7 @@ impl LogicalPlan { | PlanNode::WithCte { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Repartition { input, .. } + | PlanNode::CollectMetrics { input, .. } | PlanNode::ReplicateRows { input, .. } => p = input, PlanNode::Union { inputs, .. } => { let mut out = inputs.first()?.resolved_output()?; diff --git a/crates/krishiv-plan/src/tests/analyzer.rs b/crates/krishiv-plan/src/tests/analyzer.rs index c40a680b6..029a2a443 100644 --- a/crates/krishiv-plan/src/tests/analyzer.rs +++ b/crates/krishiv-plan/src/tests/analyzer.rs @@ -63,11 +63,11 @@ fn catalog() -> MemoryCatalog { cat } -fn ctx() -> AnalysisContext { +pub(super) fn ctx() -> AnalysisContext { context(catalog(), ResolutionConf::default()) } -fn run(plan: LogicalPlan, ctx: &AnalysisContext) -> Result { +pub(super) fn run(plan: LogicalPlan, ctx: &AnalysisContext) -> Result { let plan = Plan::try_new(plan).map_err(|error| AnalysisError { error, origin: Origin::UNKNOWN, diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index 36d8e88a4..09ded0da1 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -47,6 +47,7 @@ fn kinds_of_node(plan: &LogicalPlan, out: &mut BTreeSet<&'static str>) { PlanNode::Hint { .. } => "PlanNode::Hint", PlanNode::PipeOperator { .. } => "PlanNode::PipeOperator", PlanNode::Repartition { .. } => "PlanNode::Repartition", + PlanNode::CollectMetrics { .. } => "PlanNode::CollectMetrics", PlanNode::JoinHint { .. } => "PlanNode::JoinHint", PlanNode::Union { .. } => "PlanNode::Union", PlanNode::Distinct { .. } => "PlanNode::Distinct", @@ -236,6 +237,14 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { ), ("cast", df(1).project(vec![c("a").cast(DataType::Long)])), ("col.alias", df(1).project(vec![c("a").alias("b")])), + ( + "observe", + LogicalPlan::new(PlanNode::CollectMetrics { + input: Arc::new(df(1)), + name: "m".to_owned(), + metrics: vec![Arc::new(Expr::function("count", vec![c("a")]))], + }), + ), ] } diff --git a/crates/krishiv-plan/src/tests/mod.rs b/crates/krishiv-plan/src/tests/mod.rs index 8cc8a572a..b7da1b3f1 100644 --- a/crates/krishiv-plan/src/tests/mod.rs +++ b/crates/krishiv-plan/src/tests/mod.rs @@ -16,6 +16,7 @@ mod function_call_forms; mod inventory; mod json_tuple; mod nodes; +mod observe; mod pivot; mod prefix_cache; mod props; diff --git a/crates/krishiv-plan/src/tests/observe.rs b/crates/krishiv-plan/src/tests/observe.rs new file mode 100644 index 000000000..28241ef49 --- /dev/null +++ b/crates/krishiv-plan/src/tests/observe.rs @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `PlanNode::CollectMetrics` (`Dataset.observe`) in the analyzer: the +//! error classes and SQLSTATEs of Spark 4.2.0's `CheckAnalysis` +//! (`tests/fixtures/spark-4.2.0/observe-goldens.json`, the `err_*` cases of +//! the classic leg), the uniqueness of names across relations, and the +//! node's pass-through output. + +use super::analyzer::{ctx, run}; +use super::dsl::{c, count_star, f, int, sum}; +use crate::{Expr, ExprNode, LogicalPlan, PlanNode, WindowExpr, WindowSpec, WindowSpecRef}; +use std::sync::Arc; + +fn observe(input: LogicalPlan, name: &str, metrics: Vec) -> LogicalPlan { + LogicalPlan::new(PlanNode::CollectMetrics { + input: Arc::new(input), + name: name.to_owned(), + metrics: metrics.into_iter().map(Arc::new).collect(), + }) +} + +fn part() -> LogicalPlan { + LogicalPlan::table(["part"]) +} + +/// The error class and SQLSTATE of observing `metrics` over `part`. +fn fails(metrics: Vec) -> (String, Option) { + let e = run(observe(part(), "m", metrics), &ctx()).expect_err("must fail"); + ( + e.error.error_class().name().to_owned(), + e.error.sql_state().map(str::to_owned), + ) +} + +fn invalid(class: &str) -> (String, Option) { + (class.to_owned(), Some("42K0E".to_owned())) +} + +#[test] +fn allowed_metrics_pass_rows_through() { + let out = run( + observe( + part(), + "m", + vec![ + count_star().alias("n"), + sum(c("p_partkey")).alias("s"), + int(42), + f("sum", vec![f("rand", vec![])]).alias("r"), + f("max", vec![c("p_partkey")]), + ], + ), + &ctx(), + ) + .expect("analyzes"); + let root = out.plan.root(); + let PlanNode::CollectMetrics { metrics, input, .. } = &root.node else { + panic!("root is {root}"); + }; + assert_eq!(metrics.len(), 5); + assert_eq!( + root.resolved_output().map(|o| o.len()), + input.resolved_output().map(|o| o.len()) + ); +} + +#[test] +fn attribute_outside_an_aggregate() { + assert_eq!( + fails(vec![c("p_partkey")]), + invalid("INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE") + ); +} + +#[test] +fn nested_aggregate() { + assert_eq!( + fails(vec![sum(f("max", vec![c("p_partkey")])).alias("n")]), + invalid("INVALID_OBSERVED_METRICS.NESTED_AGGREGATES_UNSUPPORTED") + ); +} + +#[test] +fn distinct_and_filter() { + let mut call = f("count", vec![c("p_partkey")]); + if let ExprNode::Function(fc) = &mut call.node { + fc.distinct = true; + } + assert_eq!( + fails(vec![call.alias("d")]), + invalid("INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_DISTINCT_UNSUPPORTED") + ); + let mut call = f("count", vec![c("p_partkey")]); + if let ExprNode::Function(fc) = &mut call.node { + fc.filter = Some(Arc::new(f("isnotnull", vec![c("p_partkey")]))); + } + assert_eq!( + fails(vec![call.alias("f")]), + invalid("INVALID_OBSERVED_METRICS.AGGREGATE_EXPRESSION_WITH_FILTER_UNSUPPORTED") + ); +} + +#[test] +fn non_deterministic_outside_an_aggregate() { + assert_eq!( + fails(vec![f("rand", vec![]).alias("r")]), + invalid("INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC") + ); + assert_eq!( + fails(vec![ + crate::Expr::builtin(crate::Builtin::Add, vec![count_star(), f("rand", vec![])]) + .alias("r") + ]), + invalid("INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_NON_DETERMINISTIC") + ); +} + +#[test] +fn window_expression() { + let w = Expr::new(ExprNode::Window(WindowExpr { + function: Arc::new(sum(c("p_partkey"))), + window: WindowSpecRef::Spec(WindowSpec { + partition_by: vec![], + order_by: vec![], + frame: None, + }), + })); + assert_eq!( + fails(vec![w.alias("w")]), + invalid("INVALID_OBSERVED_METRICS.WINDOW_EXPRESSIONS_UNSUPPORTED") + ); +} + +#[test] +fn generator_is_an_analyzer_error() { + assert_eq!( + fails(vec![f("explode", vec![c("p_partkey")]).alias("g")]).1, + Some("42K0E".to_owned()) + ); +} + +#[test] +fn missing_name() { + let e = run(observe(part(), "", vec![count_star().alias("n")]), &ctx()).expect_err("fails"); + assert_eq!( + e.error.error_class().name(), + "INVALID_OBSERVED_METRICS.MISSING_NAME" + ); +} + +#[test] +fn unknown_column_is_unresolved() { + let e = run( + observe(part(), "m", vec![sum(c("nope")).alias("n")]), + &ctx(), + ) + .expect_err("fails"); + assert!( + e.error + .error_class() + .name() + .starts_with("UNRESOLVED_COLUMN") + ); +} + +#[test] +fn names_are_unique_across_relations_but_not_for_one_relation() { + let one = observe(part(), "dup", vec![count_star().alias("n")]); + let other = observe( + part().limit(Arc::new(int(1))), + "dup", + vec![count_star().alias("n")], + ); + let both = |l: &LogicalPlan, r: &LogicalPlan| { + l.clone().join( + r.clone(), + crate::JoinKind::Cross, + crate::JoinCondition::None, + ) + }; + let e = run(both(&one, &other), &ctx()).expect_err("two relations, one name"); + assert_eq!(e.error.error_class().name(), "DUPLICATED_METRICS_NAME"); + assert_eq!(e.error.sql_state(), Some("42710")); + // The same relation twice (a self-join) is legal. + let shared = Arc::new(one); + let twice = LogicalPlan::new(PlanNode::Join { + left: Arc::clone(&shared), + right: shared, + kind: crate::JoinKind::Cross, + condition: crate::JoinCondition::None, + }); + run(twice, &ctx()).expect("self-join of one observed relation"); +} + +#[test] +fn nested_observations_with_different_names_are_independent() { + let inner = observe(part(), "a", vec![count_star().alias("n")]); + let outer = observe(inner, "b", vec![count_star().alias("n")]); + run(outer, &ctx()).expect("analyzes"); +} diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index 6f9c97799..3147eabf4 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -398,7 +398,7 @@ impl Gen<'_> { fn plan(&mut self, depth: u8) -> LogicalPlan { let origin = self.origin(); - let kinds = if depth == 0 { 4 } else { 31 }; + let kinds = if depth == 0 { 4 } else { 32 }; let d = depth.saturating_sub(1); let node = match self.pick(kinds) { 0 => PlanNode::Read(self.source()), @@ -508,6 +508,11 @@ impl Gen<'_> { 29 => PlanNode::PipeOperator { input: self.input(d), }, + 31 => PlanNode::CollectMetrics { + input: self.input(d), + name: self.name(), + metrics: self.exprs(), + }, 30 => PlanNode::Repartition { input: self.input(d), exprs: self.exprs(), @@ -1194,6 +1199,11 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { PlanNode::Project { input, exprs } | PlanNode::Repartition { input, exprs } => { m_exprs(exprs, w) || m_plan(input, w) } + PlanNode::CollectMetrics { + input, + name, + metrics, + } => m_string(name, w) || m_exprs(metrics, w) || m_plan(input, w), PlanNode::WithColumns { input, columns } => m_columns(columns, w) || m_plan(input, w), PlanNode::Filter { input, diff --git a/crates/krishiv-plan/src/walk.rs b/crates/krishiv-plan/src/walk.rs index 2e0be19cf..e0afe4db1 100644 --- a/crates/krishiv-plan/src/walk.rs +++ b/crates/krishiv-plan/src/walk.rs @@ -200,6 +200,7 @@ pub(crate) fn all_resolved(root: Node<'_>) -> bool { | PlanNode::WithCte { .. } | PlanNode::JoinHint { .. } | PlanNode::Repartition { .. } + | PlanNode::CollectMetrics { .. } | PlanNode::Expand { .. } | PlanNode::Window { .. } | PlanNode::ReplicateRows { .. } => true, @@ -434,6 +435,11 @@ impl Teardown { | PlanNode::Repartition { input, exprs: params, + } + | PlanNode::CollectMetrics { + input, + metrics: params, + .. } => { self.plans.push(input); self.exprs.extend(params); From cc3fe1b765060b24293bcb36955341c5acf91c6c Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 00:06:59 +0000 Subject: [PATCH 03/15] exec, optimizer, plan-codec: observe operator, CollectMetrics physical node and codec feature (P3-T084 unit 3) PhysNode::CollectMetrics is plain data; the metrics lower through the global aggregate planner. The operator folds each batch into a partition-local partial aggregate state and forwards it unchanged. The optimizer keeps filters and limits above the node and keeps the metrics' inputs. Codec feature 'observe' with goldens; LIMIT 0 plans no node. The query context collects driver states and finalizes them (last instance of a name wins). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/aggregate/mod.rs | 2 + crates/krishiv-exec/src/aggregate/observe.rs | 201 ++++++++++++++++++ crates/krishiv-exec/src/context.rs | 53 +++++ crates/krishiv-exec/src/operators/project.rs | 20 +- .../krishiv-exec/src/physical/instantiate.rs | 87 +++++++- crates/krishiv-exec/src/physical/mod.rs | 30 +++ crates/krishiv-exec/src/planner/build.rs | 19 +- crates/krishiv-exec/src/planner/mod.rs | 1 + crates/krishiv-exec/src/planner/observe.rs | 64 ++++++ crates/krishiv-optimizer/src/constraints.rs | 1 + crates/krishiv-optimizer/src/dag.rs | 1 + crates/krishiv-optimizer/src/estimate.rs | 2 + crates/krishiv-optimizer/src/limit.rs | 3 +- crates/krishiv-optimizer/src/prune.rs | 3 + crates/krishiv-plan-codec/src/decode.rs | 5 + crates/krishiv-plan-codec/src/encode.rs | 15 +- crates/krishiv-plan-codec/src/features.rs | 5 + .../src/proto/krishiv.plan.v1.pb.rs | 18 +- .../tests/common/inventory.rs | 29 +++ crates/krishiv-plan-codec/tests/features.rs | 21 ++ .../tests/golden/observe/observe.bin | Bin 0 -> 893 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 3 + proto/krishiv/plan/v1/plan.proto | 12 ++ 23 files changed, 584 insertions(+), 11 deletions(-) create mode 100644 crates/krishiv-exec/src/aggregate/observe.rs create mode 100644 crates/krishiv-exec/src/planner/observe.rs create mode 100644 crates/krishiv-plan-codec/tests/golden/observe/observe.bin diff --git a/crates/krishiv-exec/src/aggregate/mod.rs b/crates/krishiv-exec/src/aggregate/mod.rs index 04dd8723b..aaf6cb809 100644 --- a/crates/krishiv-exec/src/aggregate/mod.rs +++ b/crates/krishiv-exec/src/aggregate/mod.rs @@ -71,6 +71,7 @@ mod merge_index; mod mode; mod moments; mod node; +mod observe; mod partial; mod quantile; mod rowtable; @@ -99,6 +100,7 @@ pub(crate) use collect::{collect_set_supported, kept_bytes}; pub(crate) use hll::precision_for; pub(crate) use moments::arity; pub use node::{AggregateNode, PartialNode, distinct_plan, union_distinct}; +pub(crate) use observe::{CollectMetrics, Finalizer}; pub use partial::{BypassPolicy, PartialAggregate}; pub(crate) use quantile::decimal_to_double; #[cfg(test)] diff --git a/crates/krishiv-exec/src/aggregate/observe.rs b/crates/krishiv-exec/src/aggregate/observe.rs new file mode 100644 index 000000000..c570d3275 --- /dev/null +++ b/crates/krishiv-exec/src/aggregate/observe.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `Dataset.observe` at run time ([`PhysNode::CollectMetrics`]). +//! +//! [`CollectMetrics`] is a pass-through operator: each driver folds the +//! batches it is handed into the partial states of the metrics' aggregates +//! (the table of a global partial aggregate: one group, O(1) state, weights +//! respected) and forwards every batch unchanged, so what is counted is the +//! rows that flowed through it, as Spark's `CollectMetricsExec` counts the +//! rows its consumer pulled (a limit above stops the pull; here at batch +//! granularity). A driver's states go to the query when its local state +//! ends, whether the input ended, a limit above stopped the pull or the +//! query failed (`Drop`), as the task metric of that partition; the query +//! merges the states of an observation's drivers and a [`Finalizer`] turns +//! them into one row (`QueryContext::observed`). +//! +//! Two instances of one name in a plan (a self-join of an observed +//! relation) do not add up: the last one's states are the observation's +//! (Spark's `CollectMetricsExec.collect` keeps the last). + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use krishiv_types::batch::Batch; + +use super::AggregateOutput; +use super::table::{AggCore, AggTable}; +use crate::context::TaskContext; +use crate::error::{Result, internal}; +use crate::hash::HashSeed; +use crate::operator::{LocalState, Operator, OperatorOutput}; +use crate::operators::{Project, local_mut}; + +/// Orders the instances of a plan (instantiation is one thread, in plan +/// order): the highest is the last in the plan. +static INSTANCES: AtomicU64 = AtomicU64::new(0); + +/// Turns the partial states of one observation into its row. +#[derive(Debug)] +pub(crate) struct Finalizer { + name: String, + /// This instance's place among the plan's. + instance: u64, + /// Merges the states of the drivers and gives the aggregates' results. + core: Arc, + /// The metrics over those results, one column each. + post: Arc, +} + +impl Finalizer { + /// The observation's name. + pub(crate) fn name(&self) -> &str { + &self.name + } + + /// The instance's rank in its plan. + pub(crate) fn instance(&self) -> u64 { + self.instance + } + + /// The metrics' row over `partials`, the states of every driver of the + /// instance: the aggregates' zero state if there are none. + pub(crate) fn finalize(&self, partials: &[RecordBatch]) -> Result { + let mut table = AggTable::new(&self.core, &HashSeed::new(None))?; + for states in partials { + let p = table.prepare(&Batch::new(states.clone()), true)?; + table.apply(&p)?; + } + let results = table.output_batch(&[0])?; + let out = self.post.project(&Batch::new(results), 0)?; + Ok(out.data().clone()) + } +} + +/// The operator of a [`PhysNode::CollectMetrics`](crate::physical::PhysNode). +#[derive(Debug)] +pub(crate) struct CollectMetrics { + /// The input's schema: rows pass through as they are. + schema: SchemaRef, + /// Computes the aggregates' arguments from a batch, if they are not + /// plain input columns. + pre: Option>, + /// The partial aggregate (no keys). + core: Arc, + finalizer: Arc, +} + +impl CollectMetrics { + /// The operator over rows of `schema`: `pre` computes the aggregates' + /// arguments (`pre_schema`: its output, or `schema` without it), + /// `partial` is the global aggregate's plan, `post` computes the + /// metrics from its results. + pub(crate) fn try_new( + name: &str, + schema: SchemaRef, + pre: Option>, + partial: super::AggregateSpec, + final_spec: impl FnOnce(&SchemaRef) -> Result, + post: impl FnOnce(&SchemaRef) -> Result>, + ) -> Result> { + if partial.output != AggregateOutput::Partial || !partial.group_by.is_empty() { + return Err(internal( + "observed metrics are a global partial aggregate".to_owned(), + )); + } + let core = Arc::new(AggCore::new(partial)?); + let merge = Arc::new(AggCore::new(final_spec(&core.state_schema)?)?); + let post = post(&merge.output_schema)?; + let finalizer = Arc::new(Finalizer { + name: name.to_owned(), + instance: INSTANCES.fetch_add(1, Ordering::Relaxed), + core: merge, + post, + }); + Ok(Arc::new(Self { + schema, + pre, + core, + finalizer, + })) + } +} + +/// One driver's partial states, sent to the query when it ends. +struct Local { + table: AggTable, + query: Arc, + finalizer: Arc, + partition: u32, + sent: bool, +} + +impl Local { + /// Sends the states to the query, once. + fn send(&mut self) -> Result<()> { + if std::mem::replace(&mut self.sent, true) { + return Ok(()); + } + let states = self.table.state_batch(&[0])?; + self.query.observe_partial(&self.finalizer, states); + Ok(()) + } +} + +impl Drop for Local { + fn drop(&mut self) { + // The input ended, a limit above stopped the pull, or the query + // failed: the rows seen so far count either way. + if let Err(e) = self.send() { + self.query.fail(e); + } + } +} + +impl Operator for CollectMetrics { + fn lane_safe(&self) -> bool { + true + } + + fn name(&self) -> &'static str { + "CollectMetrics" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn local_state(&self, ctx: &TaskContext) -> Result { + Ok(Box::new(Local { + table: AggTable::new(&self.core, &HashSeed::new(ctx.query().config().seed))?, + query: Arc::clone(ctx.query()), + finalizer: Arc::clone(&self.finalizer), + partition: ctx.partition(), + sent: false, + })) + } + + fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { + let l = local_mut::(local, "CollectMetrics")?; + if input.num_rows() == 0 { + return Ok(OperatorOutput::NeedMoreInput(None)); + } + let args; + let folded = match &self.pre { + Some(pre) => { + args = pre.project(input, l.partition)?; + &args + } + None => input, + }; + let p = l.table.prepare(folded, false)?; + l.table.apply(&p)?; + Ok(OperatorOutput::NeedMoreInput(Some(input.clone()))) + } + + fn finish(&self, local: &mut LocalState) -> Result> { + local_mut::(local, "CollectMetrics")?.send()?; + Ok(None) + } +} diff --git a/crates/krishiv-exec/src/context.rs b/crates/krishiv-exec/src/context.rs index 572be7208..48f43fe37 100644 --- a/crates/krishiv-exec/src/context.rs +++ b/crates/krishiv-exec/src/context.rs @@ -226,6 +226,9 @@ pub struct QueryContext { /// The scans' decode working sets (`krishiv_io::ScanGate`), made on /// first use with a budget of a quarter of the query's memory limit. scan_gate: OnceLock>, + /// The partial states of each observation (`Dataset.observe`), by name + /// and instance (see [`Self::observe_partial`]). + observed: Mutex>>, finished: AtomicBool, finished_event: Event, /// Blocking closures in flight; the query ends only once they have. @@ -233,6 +236,13 @@ pub struct QueryContext { blocking_idle: Event, } +/// One instance of an observation's node and the partial states of its +/// drivers. +struct Observed { + finalizer: Arc, + partials: Vec, +} + impl fmt::Debug for QueryContext { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("QueryContext") @@ -266,6 +276,7 @@ impl QueryContext { threads: Mutex::new(BTreeMap::new()), spills: Mutex::new(BTreeMap::new()), scan_gate: OnceLock::new(), + observed: Mutex::new(BTreeMap::new()), finished: AtomicBool::new(false), finished_event: Event::default(), blocking: AtomicUsize::new(0), @@ -273,6 +284,48 @@ impl QueryContext { }) } + /// A driver's partial states of an observation: its task metric, sent + /// when the driver's input ended (or was abandoned). The states of one + /// instance of the node merge in [`Self::observed`]. + pub(crate) fn observe_partial( + &self, + finalizer: &Arc, + states: arrow_array::RecordBatch, + ) { + let mut all = lock(&self.observed); + let instances = all.entry(finalizer.name().to_owned()).or_default(); + match instances + .iter_mut() + .find(|o| o.finalizer.instance() == finalizer.instance()) + { + Some(o) => o.partials.push(states), + None => instances.push(Observed { + finalizer: Arc::clone(finalizer), + partials: vec![states], + }), + } + } + + /// The observations (`Dataset.observe`) of the query so far: each name's + /// metrics as a one-row batch, the columns named as the metrics are. + /// An observation whose node never ran is absent. Where a name has + /// several instances in the plan (a self-join of an observed relation) + /// the last one's states are the observation's, not their sum, as in + /// Spark. + /// + /// # Errors + /// A metric's own error (finalizing runs its expressions). + pub fn observed(&self) -> Result> { + let all = lock(&self.observed); + let mut out = BTreeMap::new(); + for (name, instances) in all.iter() { + if let Some(last) = instances.iter().max_by_key(|o| o.finalizer.instance()) { + out.insert(name.clone(), last.finalizer.finalize(&last.partials)?); + } + } + Ok(out) + } + /// The query id. pub fn id(&self) -> QueryId { self.id diff --git a/crates/krishiv-exec/src/operators/project.rs b/crates/krishiv-exec/src/operators/project.rs index 3371d151b..0980cd806 100644 --- a/crates/krishiv-exec/src/operators/project.rs +++ b/crates/krishiv-exec/src/operators/project.rs @@ -149,11 +149,24 @@ impl Operator for Project { } fn execute(&self, input: &Batch, local: &mut LocalState) -> Result { - let rows = input.num_rows(); - if rows == 0 { + if input.num_rows() == 0 { return Ok(OperatorOutput::NeedMoreInput(None)); } let partition = local_mut::(local, "Project")?.partition; + Ok(OperatorOutput::NeedMoreInput(Some( + self.project(input, partition)?, + ))) + } +} + +impl Project { + /// The projection of `input`, a batch read by `partition` (no state: for + /// callers that run a projection outside a driver). + /// + /// # Errors + /// An expression's error. + pub(crate) fn project(&self, input: &Batch, partition: u32) -> Result { + let rows = input.num_rows(); let data = input.data(); let columns = self .items @@ -174,7 +187,6 @@ impl Operator for Project { let options = RecordBatchOptions::new().with_row_count(Some(rows)); let out = RecordBatch::try_new_with_options(Arc::clone(&self.schema), columns, &options) .map_err(|e| internal(format!("project produced an invalid batch: {e}")))?; - let out = Batch::with_weights(out, input.weights().clone())?.with_split(input.split()); - Ok(OperatorOutput::NeedMoreInput(Some(out))) + Ok(Batch::with_weights(out, input.weights().clone())?.with_split(input.split())) } } diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index cbb7f9136..f1bc9918e 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -57,8 +57,8 @@ use super::{ SortMergeJoinPlan, SubqueryPlan, ValueId, Version, WritePlan, files_of, partitions_of, }; use crate::aggregate::{ - AggInput, AggregateExpr, AggregateNode, AggregateOutput, AggregateSpec, HashAggregate, - distinct_plan, + AggInput, AggregateExpr, AggregateNode, AggregateOutput, AggregateSpec, CollectMetrics, + HashAggregate, distinct_plan, }; use crate::error::{Result, internal}; use crate::exchange::{Exchange, ExchangeRead, ExchangeSink, ExchangeWrite, MapOutputId}; @@ -1035,6 +1035,13 @@ impl Instantiation<'_> { } => self.with_subqueries(input, subqueries, |input, v| { Ok(project_op(input, layout, items, nullable.as_deref(), v)? as Arc) }), + PhysNode::CollectMetrics { + input, + name, + measure, + } => self.with_subqueries(input, &[], |input, v| { + collect_metrics_op(name, measure, input, v) + }), PhysNode::ReplicateRows { input, layout, @@ -1693,6 +1700,82 @@ pub(crate) fn expand_op( Ok(Expand::try_new(projects)?) } +/// The operator of a `CollectMetrics` node over rows of `schema`: the +/// shape of `measure` is what `planner/observe.rs` makes. +pub(crate) fn collect_metrics_op( + name: &str, + measure: &PhysNode, + schema: &SchemaRef, + v: &Values<'_>, +) -> Result> { + let bad = || { + internal(format!( + "the metrics of observation {name} have an unexpected plan" + )) + }; + let PhysNode::Project { + input: aggregate, + layout: post_layout, + items: post_items, + nullable: post_nullable, + subqueries, + } = measure + else { + return Err(bad()); + }; + let PhysNode::Aggregate { + input: args, + aggregate: plan, + } = aggregate.as_ref() + else { + return Err(bad()); + }; + if !subqueries.is_empty() + || !plan.group_by.is_empty() + || plan.output != AggregateOutput::Final + || plan.distinct.iter().any(|d| *d) + { + return Err(bad()); + } + let (pre, pre_schema) = match args.as_ref() { + PhysNode::OneRow => (None, Arc::clone(schema)), + PhysNode::Project { + input: leaf, + layout, + items, + nullable, + subqueries, + } if matches!(leaf.as_ref(), PhysNode::OneRow) && subqueries.is_empty() => { + let p = project_op(schema, layout, items, nullable.as_deref(), v)?; + let s = p.schema(); + (Some(p), s) + } + _ => return Err(bad()), + }; + let partial = AggregatePlan { + output: AggregateOutput::Partial, + ..plan.clone() + }; + let merge = final_of(&partial)?; + CollectMetrics::try_new( + name, + Arc::clone(schema), + pre, + aggregate_spec(&pre_schema, &partial, v.env, None), + |states| Ok(aggregate_spec(states, &merge, v.env, None)), + |results| { + project_op( + results, + post_layout, + post_items, + post_nullable.as_deref(), + v, + ) + }, + ) + .map(|op| op as Arc) +} + pub(crate) fn replicate_op( schema: &SchemaRef, layout: &[AttrId], diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index ccc10544c..584e6abf6 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -421,6 +421,22 @@ pub enum PhysNode { /// The calls, in output column order. calls: Vec, }, + /// `Dataset.observe`: the rows of `input` pass through unchanged while + /// each partition folds them into the metrics' aggregate states, which + /// the query's result reports under `name`. `measure` is plain data, + /// not a child: `Project(Aggregate(Project(OneRow)))` with `OneRow` + /// standing for the input's rows, as the planner lowers the metrics + /// as a global aggregate (the first projection computes the aggregates' + /// arguments, the aggregate is the partial phase's plan, the last + /// projection computes the metrics from its results). + CollectMetrics { + /// The input. + input: Box, + /// The observation's name. + name: String, + /// The metrics' computation over the input's columns. + measure: Box, + }, /// Each row of `input` `count` times (`INTERSECT ALL`, `EXCEPT ALL`). ReplicateRows { /// The input. @@ -1021,6 +1037,7 @@ plan_data!(enum PhysNode { Filter { input, layout, predicate, subqueries }, Project { input, layout, items, nullable, subqueries }, Limit { input, offset, limit }, + CollectMetrics { input, name, measure }, ReplicateRows { input, layout, count, subqueries }, PythonEval { input, layout, calls }, Generate { input, layout, kind, outer, arg, subqueries }, @@ -1127,6 +1144,7 @@ impl PhysNode { | Self::WindowGroupLimit { input, .. } | Self::Aggregate { input, .. } | Self::PythonEval { input, .. } + | Self::CollectMetrics { input, .. } | Self::WriteFiles { input, .. } => vec![input.as_ref()], Self::Exchange(e) => vec![e.input.as_ref()], Self::HashJoin(j) => vec![j.probe.as_ref(), j.build.as_ref()], @@ -1139,6 +1157,17 @@ impl PhysNode { } } + /// The plan of an observation's metrics (`CollectMetrics`): plain data + /// that is not a child (it reads no relation), listed for the codec and + /// the feature scan. + #[must_use] + pub fn measure(&self) -> Option<&PhysNode> { + match self { + Self::CollectMetrics { measure, .. } => Some(measure), + _ => None, + } + } + /// [`Self::children`], mutable. pub fn children_mut(&mut self) -> Vec<&mut PhysNode> { match self { @@ -1173,6 +1202,7 @@ impl PhysNode { | Self::WindowGroupLimit { input, .. } | Self::Aggregate { input, .. } | Self::PythonEval { input, .. } + | Self::CollectMetrics { input, .. } | Self::WriteFiles { input, .. } => vec![input.as_mut()], Self::Exchange(e) => vec![e.input.as_mut()], Self::HashJoin(j) => vec![j.probe.as_mut(), j.build.as_mut()], diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index 40f78f056..272252594 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -188,6 +188,7 @@ fn kept_order(node: &PhysNode, order: &[OrderKey], layout: &[AttrId]) -> Vec order .iter() @@ -308,6 +309,9 @@ pub(crate) struct Planner<'a> { loops: HashMap, /// Build-side projections requested and applied (`build_project`). pub(crate) moves: super::build_project::Moves, + /// Nesting depth of the relations whose rows are never read (under a + /// `LIMIT 0`): `observe` plans no node there. + pub(super) unobserved: u32, } impl<'a> Planner<'a> { @@ -342,6 +346,7 @@ impl<'a> Planner<'a> { scan_filters: HashMap::new(), loops: HashMap::new(), moves: super::build_project::Moves::default(), + unobserved: 0, } } @@ -523,14 +528,24 @@ impl<'a> Planner<'a> { } PlanNode::Limit { input, limit } => { let n = limit_value(limit)?; + // Spark replaces what a `LIMIT 0` reads with an empty + // relation, so an `observe` under it reports nothing. + self.unobserved += u32::from(n == 0); // `LIMIT n OFFSET m` is the limit over the offset: skip, then keep. - match &input.node { + let planned = match &input.node { PlanNode::Offset { input, offset } => { self.limit(input, limit_value(offset)?, n) } _ => self.limit(input, 0, n), - } + }; + self.unobserved -= u32::from(n == 0); + planned } + PlanNode::CollectMetrics { + input, + name, + metrics, + } => self.collect_metrics(input, name, metrics), PlanNode::Offset { input, offset } => self.limit(input, limit_value(offset)?, u64::MAX), PlanNode::Sort { input, diff --git a/crates/krishiv-exec/src/planner/mod.rs b/crates/krishiv-exec/src/planner/mod.rs index 6c3e67b0c..7d007da47 100644 --- a/crates/krishiv-exec/src/planner/mod.rs +++ b/crates/krishiv-exec/src/planner/mod.rs @@ -35,6 +35,7 @@ mod join; mod late_fetch; mod ndv_sample; pub(crate) mod nodes; +mod observe; mod orient; mod probe_select; mod properties; diff --git a/crates/krishiv-exec/src/planner/observe.rs b/crates/krishiv-exec/src/planner/observe.rs new file mode 100644 index 000000000..22a7496f3 --- /dev/null +++ b/crates/krishiv-exec/src/planner/observe.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `Dataset.observe`: [`PhysNode::CollectMetrics`]. +//! +//! The metrics are a global aggregate over the rows that pass through, so +//! they plan as one (`Planner::aggregate` over a placeholder relation +//! with the input's columns), and the node keeps the resulting +//! projections and aggregate as plain data (`measure`). The rows +//! themselves are planned as they are: nothing moves below the node (a +//! runtime filter would change what the metrics count) and nothing above +//! it sees a difference. + +use std::sync::Arc; + +use krishiv_plan::{Expr, Grouping, LogicalPlan}; + +use super::build::{Planned, Planner, internal}; +use super::explain::ExplainNode; +use super::properties::Placement; +use crate::error::Result; +use crate::physical::PhysNode; + +impl Planner<'_> { + /// `input`'s rows, observed as `metrics` under `name`. Where the rows + /// are never read (a `LIMIT 0` above, which Spark replaces with an + /// empty relation) the node is not planned. + pub(super) fn collect_metrics( + &mut self, + input: &Arc, + name: &str, + metrics: &[Arc], + ) -> Result { + let planned = self.plan(input)?; + if self.unobserved > 0 { + return Ok(planned); + } + let rows = Planned { + node: PhysNode::OneRow, + schema: Arc::clone(&planned.schema), + layout: planned.layout.clone(), + explain: ExplainNode::new("ObservedRows", String::new(), Vec::new()), + sites: Vec::new(), + placement: Placement::Single, + ordering: Vec::new(), + }; + let measured = self.aggregate(rows, &Grouping::GroupBy(Vec::new()), metrics)?; + if !matches!(measured.node, PhysNode::Project { .. }) { + return Err(internal("observed metrics did not plan as a projection")); + } + let (schema, layout) = (Arc::clone(&planned.schema), planned.layout.clone()); + let mut out = planned.wrap( + |input| PhysNode::CollectMetrics { + input, + name: name.to_owned(), + measure: Box::new(measured.node), + }, + schema, + layout, + ("CollectMetrics", name.to_owned()), + ); + // A filter below would change which rows the metrics count. + out.sites.clear(); + Ok(out) + } +} diff --git a/crates/krishiv-optimizer/src/constraints.rs b/crates/krishiv-optimizer/src/constraints.rs index e1901a01e..91bd4d43a 100644 --- a/crates/krishiv-optimizer/src/constraints.rs +++ b/crates/krishiv-optimizer/src/constraints.rs @@ -697,6 +697,7 @@ impl Memo { | PlanNode::WithCte { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Repartition { input, .. } + | PlanNode::CollectMetrics { input, .. } | PlanNode::ReplicateRows { input, .. } | PlanNode::SubqueryAlias { input, .. } if !matches!( diff --git a/crates/krishiv-optimizer/src/dag.rs b/crates/krishiv-optimizer/src/dag.rs index 38333ae36..ea2c12eab 100644 --- a/crates/krishiv-optimizer/src/dag.rs +++ b/crates/krishiv-optimizer/src/dag.rs @@ -121,6 +121,7 @@ pub(crate) fn with_input(node: &Arc, input: Arc) -> Ar | PlanNode::Shift { input: slot, .. } | PlanNode::JoinHint { input: slot, .. } | PlanNode::Repartition { input: slot, .. } + | PlanNode::CollectMetrics { input: slot, .. } | PlanNode::Window { input: slot, .. } | PlanNode::Generate { input: slot, .. } | PlanNode::ReplicateRows { input: slot, .. } => *slot = input, diff --git a/crates/krishiv-optimizer/src/estimate.rs b/crates/krishiv-optimizer/src/estimate.rs index c9d100e30..ece06aa6e 100644 --- a/crates/krishiv-optimizer/src/estimate.rs +++ b/crates/krishiv-optimizer/src/estimate.rs @@ -417,6 +417,8 @@ impl<'a> Estimator<'a> { | PlanNode::SubqueryAlias { input, .. } | PlanNode::JoinHint { input, .. } | PlanNode::Repartition { input, .. } + // Rows pass through unchanged. + | PlanNode::CollectMetrics { input, .. } // A window adds columns (without statistics), never rows. | PlanNode::Window { input, .. } // Each distinct row of an `INTERSECT ALL` / `EXCEPT ALL` count: diff --git a/crates/krishiv-optimizer/src/limit.rs b/crates/krishiv-optimizer/src/limit.rs index 461f11a72..b9582155a 100644 --- a/crates/krishiv-optimizer/src/limit.rs +++ b/crates/krishiv-optimizer/src/limit.rs @@ -72,7 +72,8 @@ pub(crate) fn max_rows(plan: &LogicalPlan) -> Option { | PlanNode::Shift { input, .. } | PlanNode::Window { input, .. } | PlanNode::JoinHint { input, .. } - | PlanNode::Repartition { input, .. } => max_rows(input), + | PlanNode::Repartition { input, .. } + | PlanNode::CollectMetrics { input, .. } => max_rows(input), PlanNode::Union { inputs, .. } => inputs .iter() .try_fold(0_u64, |n, i| n.checked_add(max_rows(i)?)), diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index 708558638..f16b21012 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -125,6 +125,9 @@ impl Graph { // A repartition reads its expressions' columns whatever is // read above it: Spark evaluates them under `count(*)` too. | PlanNode::Repartition { input, .. } + // The metrics read their columns whatever is read above: a + // `count` action prunes everything else, not these. + | PlanNode::CollectMetrics { input, .. } | PlanNode::Generate { input, .. } | PlanNode::ReplicateRows { input, .. } => { let mut need = req.clone(); diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index ca2544fc0..b3a9e8335 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -532,6 +532,11 @@ impl Decoder { offset: l.offset, limit: l.limit, }, + Kind::CollectMetrics(m) => PhysNode::CollectMetrics { + input: self.boxed(m.input)?, + name: m.name.clone(), + measure: self.boxed(m.measure)?, + }, Kind::ReplicateRows(r) => PhysNode::ReplicateRows { input: self.boxed(r.input)?, layout: layout(&r.layout), diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 485d897a8..49c77cadb 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -56,7 +56,11 @@ pub fn to_message(plan: &PhysicalPlan) -> Result { let features = plan_features(plan); let mut e = Encoder::default(); for tree in plan.roots() { - let order = post_order(tree, PhysNode::children); + let order = post_order(tree, |n| { + let mut v = n.children(); + v.extend(n.measure()); + v + }); for node in &order { let message = e.node(node)?; e.nodes.insert(std::ptr::from_ref(*node), e.table.len()); @@ -306,6 +310,15 @@ impl Encoder { limit: *limit, offset: *offset, }), + PhysNode::CollectMetrics { + input, + name, + measure, + } => Kind::CollectMetrics(pb::CollectMetricsNode { + input: self.node_index(input)?, + name: name.clone(), + measure: self.node_index(measure)?, + }), PhysNode::ReplicateRows { input, layout: l, diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index 189160f9d..c19fff0b9 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -149,6 +149,7 @@ registry! { (mode_regr_bit_listagg_aggs, ModeRegrBitListaggAggs, POST_MIGRATION_READER_VERSION, Semantic, "The `mode`, regression (`regr_*`), bitwise (`bit_and`, `bit_or`, `bit_xor`), `listagg`/`string_agg` and `try_sum`/`try_avg` aggregates."), (moment_aggregates, MomentAggregates, POST_MIGRATION_READER_VERSION, Semantic, "The moment aggregates (`stddev_samp`, `stddev_pop`, `var_samp`, `var_pop`, `skewness`, `kurtosis`, `corr`, `covar_samp`, `covar_pop`)."), (nested_types, NestedTypes, POST_MIGRATION_READER_VERSION, Semantic, "Array, map and struct types (`DataType.KIND_NESTED`, written as Spark's JSON form of the type)."), + (observe, Observe, POST_MIGRATION_READER_VERSION, Semantic, "The `CollectMetrics` node (`Dataset.observe`): rows pass through while the metrics of a named observation are folded; a worker without it cannot report them."), (options_map_literals, OptionsMapLiterals, POST_MIGRATION_READER_VERSION, Semantic, "A `MAP` literal: the options of one call (`from_json`, `from_csv`, `from_xml`). Entries are classified as source options are: a sensitive key or a URL with credentials does not encode (ADR 0038)."), (partial_aggregate, PartialAggregate, POST_MIGRATION_READER_VERSION, Optimization, "Partial aggregates below a join, whose states a later aggregate merges (eager aggregation)."), (partitioned_writes, PartitionedWrites, POST_MIGRATION_READER_VERSION, Semantic, "A `WritePlan` laid out by partition columns (Spark's directory layout)."), @@ -412,6 +413,7 @@ pub fn feature_of(construct: Construct<'_>) -> Option { PhysNode::SortMergeJoin(_) => Some(Feature::SortMergeJoin), PhysNode::WindowGroupLimit { .. } => Some(Feature::WindowGroupLimit), PhysNode::PythonEval { .. } => Some(Feature::PythonUdf), + PhysNode::CollectMetrics { .. } => Some(Feature::Observe), // A partitioned write needs a reader that // lays files out by partition (the older fields would be // ignored and the rows written flat). @@ -708,6 +710,7 @@ fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { | PhysNode::TopN { .. } | PhysNode::WindowGroupLimit { .. } | PhysNode::PythonEval { .. } + | PhysNode::CollectMetrics { .. } | PhysNode::Union { .. } => {} PhysNode::WriteFiles { write, .. } => { if write.as_ref().is_some_and(|w| !w.column_types.is_empty()) { @@ -851,6 +854,7 @@ fn exprs_of<'a>(node: &'a PhysNode) -> Vec<&'a Arc> { | PhysNode::Window { .. } | PhysNode::WindowGroupLimit { .. } | PhysNode::Aggregate { .. } + | PhysNode::CollectMetrics { .. } | PhysNode::Union { .. } | PhysNode::WriteFiles { .. } | PhysNode::Exchange(_) @@ -925,6 +929,7 @@ pub fn plan_features(plan: &PhysicalPlan) -> BTreeSet { out.extend(feature_of(c)); }); stack.extend(node.children()); + stack.extend(node.measure()); } for field in plan.schema().fields() { visit_arrow_field(field, &mut |c| out.extend(feature_of(c))); diff --git a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index 685b2ba94..5f3cf79d3 100644 --- a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs +++ b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs @@ -145,7 +145,7 @@ pub struct ExplainNode { pub struct Node { #[prost( oneof = "node::Kind", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 21, 24, 25, 26" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 21, 24, 25, 26, 27" )] pub kind: ::core::option::Option, } @@ -217,8 +217,24 @@ pub mod node { /// feature: python_udf #[prost(message, tag = "26")] PythonEval(super::PythonEvalNode), + /// feature: observe + #[prost(message, tag = "27")] + CollectMetrics(super::CollectMetricsNode), } } +/// feature: observe +/// `Dataset.observe`: the rows of `input` pass through unchanged while the +/// metrics (`measure`, a plan over no relation: `OneRow` stands for the input's +/// rows) are folded into an observation named `name`. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct CollectMetricsNode { + #[prost(uint32, tag = "1")] + pub input: u32, + #[prost(string, tag = "2")] + pub name: ::prost::alloc::string::String, + #[prost(uint32, tag = "3")] + pub measure: u32, +} /// feature: python_udf /// The rows of `input`, then one column per Python UDF call over every row. #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index f86b3921a..852e0b002 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -1693,6 +1693,34 @@ pub(crate) fn udf_call() -> PhysNode { } } +/// `Dataset.observe` (`observe`): a scan whose rows pass through a +/// `CollectMetrics` whose measure is `Project(Aggregate(OneRow))`. +pub(crate) fn collect_metrics() -> PhysNode { + let aggregates = vec![aggregate_exprs().remove(0)]; + PhysNode::CollectMetrics { + input: boxed(scan()), + name: "obs".to_owned(), + measure: boxed(PhysNode::Project { + input: boxed(PhysNode::Aggregate { + input: boxed(PhysNode::OneRow), + aggregate: AggregatePlan { + group_by: Vec::new(), + distinct: vec![false], + aggregates, + output: AggregateOutput::Final, + }, + }), + layout: vec![AttrId(9)], + items: vec![ProjectItemPlan::Column { + index: 0, + name: Some("n".to_owned()), + }], + nullable: None, + subqueries: Vec::new(), + }), + } +} + /// A Python UDF call (`python_udf`): a `PythonEval` appending the call's /// column, read by the projection above it. The call references the UDF by /// name and registration version only. @@ -1916,6 +1944,7 @@ pub(crate) fn cases() -> Vec<(String, PhysicalPlan)> { out.push(("expressions".to_owned(), plan_of(expressions()))); out.push(("rust_udf".to_owned(), plan_of(udf_call()))); out.push(("python_udf".to_owned(), plan_of(python_eval()))); + out.push(("observe".to_owned(), plan_of(collect_metrics()))); out.push(("rust_udaf".to_owned(), plan_of(udaf_call()))); // The expression environment flag the `legacy_statistical_aggregate` // feature is about: the same root as `one_row`, under a session with diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs index 85d10202b..c1931d390 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -2100,3 +2100,24 @@ fn a_python_udf_needs_its_feature() { } ); } + +/// A plan observing metrics lists `observe`; a reader without it refuses +/// the plan naming the feature. +#[test] +fn an_observation_needs_its_feature() { + assert_eq!(message("observe").features, ["observe"]); + assert_eq!(Feature::Observe.kind(), FeatureKind::Semantic); + let bytes = message("observe").encode_to_vec(); + let older = ReaderFeatures::only( + Feature::ALL + .iter() + .copied() + .filter(|f| *f != Feature::Observe), + ); + assert_eq!( + decode_with(&bytes, &Limits::default(), &older).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["observe".to_owned()] + } + ); +} diff --git a/crates/krishiv-plan-codec/tests/golden/observe/observe.bin b/crates/krishiv-plan-codec/tests/golden/observe/observe.bin new file mode 100644 index 0000000000000000000000000000000000000000..cce7548ac3eb9c727e73b123dbf23f430daabbf5 GIT binary patch literal 893 zcmbVKv2N5r5ZzhZ=e@}#arS~RD+=qX$P&#}CMh0X~5OiLc-bQt%IE_8cNbq{^(`XlG_;-p=fbAR4H1j-R*)<#T+7>+Pv-PF^2t ztE+mBVP4Giva$$PM(Dq-^Fj(7w6-aat(;(B>yyIBd-5*&r%j#9%`@E;i+b4>4^|LKV95n2pNV~S=fkOPnW7xl~E_!~m$)o@LgIbvZeq zA3FPQ=sh=!Z)&=|%DeqzX{s?DqW`p*SqW9Y)iVk2r-(sT&KFhzOvmVFb!{bmT&v59 z{u<*Cy?v>ZMQLnurrV^Po4P6TgIGY+{q*9GpSwRlzJD7`LTz;Qp|#oiwO^a5G{#}s zn%bIWRmDJ$x61|JgC~&rlJ9J1<8%w*5k-jYDs(3uqo+;o+#dXMzX<7li{@Q%1b+b| Ch-(7? literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index b1dfbc008..618bcdfc7 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -45,6 +45,7 @@ variants!( ReplicateRows, Generate, PythonEval, + CollectMetrics, Expand, Sort, TopN, @@ -75,6 +76,7 @@ fn variant_of(node: &PhysNode) -> Variant { PhysNode::ReplicateRows { .. } => Variant::ReplicateRows, PhysNode::Generate { .. } => Variant::Generate, PhysNode::PythonEval { .. } => Variant::PythonEval, + PhysNode::CollectMetrics { .. } => Variant::CollectMetrics, PhysNode::UnionLoop { .. } => Variant::UnionLoop, PhysNode::LoopRef { .. } => Variant::LoopRef, PhysNode::Expand { .. } => Variant::Expand, @@ -620,6 +622,7 @@ fn readers_that_predate_features_still_refuse_what_they_cannot_read() { | "session_zone_name" | "aggregate_filter" | "range" + | "observe" | "aggregate_collects" | "aggregate_vectors" | "aggregate_mode_regr_bit_listagg" diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 3113d37e9..9634eee38 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -172,9 +172,21 @@ message Node { ExchangeSourceNode exchange_source = 25; // feature: python_udf PythonEvalNode python_eval = 26; + // feature: observe + CollectMetricsNode collect_metrics = 27; } } +// feature: observe +// `Dataset.observe`: the rows of `input` pass through unchanged while the +// metrics (`measure`, a plan over no relation: `OneRow` stands for the input's +// rows) are folded into an observation named `name`. +message CollectMetricsNode { + uint32 input = 1; + string name = 2; + uint32 measure = 3; +} + // feature: python_udf // The rows of `input`, then one column per Python UDF call over every row. message PythonEvalNode { From 00afb46a599d8dbe8e6fc069f039f78317f5356e Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 00:09:45 +0000 Subject: [PATCH 04/15] session: DataFrame::observe and Observation (P3-T084 unit 4) Observation handle (get blocks, first action wins, completes on failure and cancellation), delivery from the query context when the result stream ends, KRISHIV_OBSERVATION_REUSED and KRISHIV_NO_OBSERVE_BEFORE_GET, and end-to-end tests over the Spark 4.2.0 goldens, parallelism exactness and cancellation. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-common/krishiv-error-classes.json | 12 + crates/krishiv-session/src/catalog.rs | 1 + crates/krishiv-session/src/dataframe.rs | 1 + .../krishiv-session/src/dataframe/observe.rs | 124 ++++++++ .../krishiv-session/src/dataframe/traced.rs | 5 +- crates/krishiv-session/src/lib.rs | 2 + crates/krishiv-session/src/observation.rs | 224 ++++++++++++++ crates/krishiv-session/tests/observe.rs | 275 ++++++++++++++++++ 8 files changed, 642 insertions(+), 2 deletions(-) create mode 100644 crates/krishiv-session/src/dataframe/observe.rs create mode 100644 crates/krishiv-session/src/observation.rs create mode 100644 crates/krishiv-session/tests/observe.rs diff --git a/crates/krishiv-common/krishiv-error-classes.json b/crates/krishiv-common/krishiv-error-classes.json index e060f46d7..4382d7e22 100644 --- a/crates/krishiv-common/krishiv-error-classes.json +++ b/crates/krishiv-common/krishiv-error-classes.json @@ -110,6 +110,12 @@ ], "sqlState" : "42501" }, + "KRISHIV_NO_OBSERVE_BEFORE_GET" : { + "message" : [ + "The observation was read before a DataFrame was observed with it. Call `DataFrame.observe` with it first. (PySpark's `NO_OBSERVE_BEFORE_GET`.)" + ], + "sqlState" : "42K0E" + }, "KRISHIV_OBJECT_STORE_REQUEST_FAILED" : { "message" : [ "Reading the file failed: the object store did not complete the request after attempt(s): . This is a storage failure, not a corrupt file, so it is never skipped by \"ignoreCorruptFiles\". Retry the query, or check the store's availability and permissions." @@ -117,6 +123,12 @@ "sqlState" : "58030", "audience" : "operator" }, + "KRISHIV_OBSERVATION_REUSED" : { + "message" : [ + "The observation already observes a DataFrame: an Observation can be used with a DataFrame only once. (PySpark's `REUSE_OBSERVATION`.)" + ], + "sqlState" : "42K0E" + }, "KRISHIV_OPERATION_LIMIT" : { "message" : [ "Cannot start a new operation: allows at most live operations per session. Wait for one to finish or release it (ReleaseExecute, or interrupt it in the client), or raise \"\"." diff --git a/crates/krishiv-session/src/catalog.rs b/crates/krishiv-session/src/catalog.rs index 657116090..ed63b7752 100644 --- a/crates/krishiv-session/src/catalog.rs +++ b/crates/krishiv-session/src/catalog.rs @@ -950,6 +950,7 @@ fn streaming_leaves(root: &PhysNode) -> (Vec<&krishiv_exec::physical::ScanNode>, PhysNode::Filter { input, .. } | PhysNode::Project { input, .. } | PhysNode::ReplicateRows { input, .. } + | PhysNode::CollectMetrics { input, .. } | PhysNode::Generate { input, .. } | PhysNode::Expand { input, .. } | PhysNode::Limit { input, .. } diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index 1a98c6f72..321899562 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -31,6 +31,7 @@ use crate::session::{QueryEnv, RecordBatchStream, Session, WriteTarget, read_fil pub(crate) mod checkpoint; mod fill_stream; mod geo_gate; +mod observe; mod traced; pub(crate) use fill_stream::{CacheFill, FillKind}; diff --git a/crates/krishiv-session/src/dataframe/observe.rs b/crates/krishiv-session/src/dataframe/observe.rs new file mode 100644 index 000000000..24f605250 --- /dev/null +++ b/crates/krishiv-session/src/dataframe/observe.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `DataFrame::observe` and the delivery of an observed plan's metrics to +//! its [`Observation`] handles when an action's result stream ends (P3-T084). +//! A child of `dataframe` so it reads the frame's private fields. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow_array::RecordBatch; +use futures::Stream; +use krishiv_common::error::KrishivError; +use krishiv_exec::QueryContext; +use krishiv_plan::PlanNode; + +use super::DataFrame; +use crate::column::Column; +use crate::observation::{Observation, Shared}; +use crate::session::RecordBatchStream; + +impl DataFrame { + /// This frame, observed: `metrics` (aggregate expressions, evaluated as a + /// global aggregate over the rows that flow through; unaliased ones are + /// named as Spark names them) are collected as the rows pass, and + /// `observation` returns them once the first action has finished + /// (Spark's `Dataset.observe(Observation, ...)`). + /// + /// # Errors + /// `KRISHIV_OBSERVATION_REUSED` when `observation` already observes a + /// frame. Invalid metrics (a column outside an aggregate, `DISTINCT`, a + /// window, ...) fail the first action, as analysis errors do. + pub fn observe( + &self, + observation: &Observation, + metrics: impl IntoIterator>, + ) -> Result { + observation.shared.bind()?; + let mut out = self.observe_named(observation.name(), metrics); + out.keep.push(Arc::clone(&observation.shared) as _); + Ok(out) + } + + /// As [`Self::observe`] with a name and no handle (Spark's + /// `Dataset.observe(String, ...)`): the metrics are computed with the + /// query and reported to whoever reads the query's observations. + #[must_use] + pub fn observe_named( + &self, + name: &str, + metrics: impl IntoIterator>, + ) -> Self { + self.derive(PlanNode::CollectMetrics { + input: Arc::clone(&self.plan), + name: name.to_owned(), + metrics: metrics + .into_iter() + .map(|c| Arc::new(c.into().expr)) + .collect(), + }) + } + + /// The observation handles this frame carries. + pub(super) fn observations(&self) -> Vec> { + self.keep + .iter() + .filter_map(|k| Arc::clone(k).downcast::().ok()) + .collect() + } +} + +/// `stream`, completing `observers` from `query`'s observations when it ends, +/// fails or is dropped (Spark completes them whether the action succeeded or +/// not); the first action wins. +pub(super) fn deliver( + stream: RecordBatchStream, + query: Arc, + observers: Vec>, +) -> RecordBatchStream { + if observers.is_empty() { + return stream; + } + Box::pin(Delivering { + inner: stream, + query, + observers, + }) +} + +struct Delivering { + inner: RecordBatchStream, + query: Arc, + observers: Vec>, +} + +impl Delivering { + fn complete(&mut self) { + let observers = std::mem::take(&mut self.observers); + if observers.is_empty() { + return; + } + let observed = self.query.observed().unwrap_or_default(); + for o in observers { + o.complete(observed.get(o.name())); + } + } +} + +impl Stream for Delivering { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let polled = self.inner.as_mut().poll_next(cx); + if matches!(polled, Poll::Ready(None | Some(Err(_)))) { + self.complete(); + } + polled + } +} + +impl Drop for Delivering { + fn drop(&mut self) { + self.complete(); + } +} diff --git a/crates/krishiv-session/src/dataframe/traced.rs b/crates/krishiv-session/src/dataframe/traced.rs index 647b1c215..ab8e3c70f 100644 --- a/crates/krishiv-session/src/dataframe/traced.rs +++ b/crates/krishiv-session/src/dataframe/traced.rs @@ -69,11 +69,12 @@ impl DataFrame { // The execution span stays open while the rows are produced and read: it closes when // the result stream is dropped, not when `run` returns the stream. let exec = tracing::info_span!("krishiv.execute"); - let stream = self + let (stream, query) = self .session - .run(&physical, &a.env) + .run_with_query(&physical, &a.env) .instrument(exec.clone()) .await?; + let stream = super::observe::deliver(stream, query, self.observations()); let stream = crate::telemetry::hold_span(stream, exec); let stream = fill.fill(stream, &self.session.state.cache); let stream = match result { diff --git a/crates/krishiv-session/src/lib.rs b/crates/krishiv-session/src/lib.rs index aaf5ac9c7..2c880ec0f 100644 --- a/crates/krishiv-session/src/lib.rs +++ b/crates/krishiv-session/src/lib.rs @@ -69,6 +69,7 @@ pub mod logging; mod memory_table; mod monitor; mod nested_value; +mod observation; #[cfg(feature = "telemetry-otlp")] pub mod otlp; mod partitions; @@ -130,6 +131,7 @@ pub use krishiv_udf_sdk::{ ReturnField, ScalarUdf, Signature, StateField, UdfArg, UdfError, UdfInputType, UdfLibrary, UdfProperties, }; +pub use observation::Observation; pub use profile::{ExplainVerbose, OptimizerStep, VERBOSE_FORMAT, VERBOSE_VERSION}; pub use python_udfs::{ANONYMOUS_VERSION_BIT, PythonUdfDef, PythonUdfEntry, PythonUdfLibrary}; pub use runtime::{EngineRuntime, EngineRuntimeBuilder}; diff --git a/crates/krishiv-session/src/observation.rs b/crates/krishiv-session/src/observation.rs new file mode 100644 index 000000000..7bdc52daa --- /dev/null +++ b/crates/krishiv-session/src/observation.rs @@ -0,0 +1,224 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `Observation` (Spark's `org.apache.spark.sql.Observation`, PySpark's +//! `pyspark.sql.Observation`): the handle `DataFrame::observe` binds to one +//! frame, whose [`Observation::get`] returns the metrics of the first +//! action that ran the observed plan (P3-T084; research +//! `t084-observe-semantics.md` section 3). +//! +//! - One handle observes one frame: a second `observe` with it fails +//! (`KRISHIV_OBSERVATION_REUSED`; PySpark's `REUSE_OBSERVATION`). +//! - [`Observation::get`] blocks until an action has finished (on success or +//! on failure, as Spark completes the promise in a `finally`) and never +//! again changes: the first action wins. Without any action it blocks +//! forever, as Spark's does; [`Observation::get_timeout`] bounds the wait. +//! Before `observe` bound the handle it fails +//! (`KRISHIV_NO_OBSERVE_BEFORE_GET`; PySpark's `NO_OBSERVE_BEFORE_GET`). +//! - The value is one row, a [`RecordBatch`] of one row whose columns are +//! the metrics, named as written (`count(1)`, `max(v)` when not aliased); +//! where the observed plan never ran its node (a `LIMIT 0` above it) the +//! row has no columns (Spark's `{}`). + +use std::sync::{Arc, Condvar, Mutex, PoisonError}; +use std::time::Duration; + +use arrow_array::RecordBatch; +use arrow_schema::Schema; +use krishiv_common::error::{ + InvalidObservedMetricsMissingName, KrishivError, KrishivNoObserveBeforeGet, + KrishivObservationReused, +}; + +/// Where an observation is. +enum State { + /// Not bound to a frame by `observe` yet. + Unbound, + /// Bound; no action has finished. + Waiting, + /// The first action finished. + Done(RecordBatch), +} + +/// What the frames and the handle share. +pub(crate) struct Shared { + name: String, + state: Mutex, + done: Condvar, +} + +impl Shared { + fn lock(&self) -> std::sync::MutexGuard<'_, State> { + self.state.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// The name keying the plan's `CollectMetrics` node. + pub(crate) fn name(&self) -> &str { + &self.name + } + + /// Binds the handle to a frame. + pub(crate) fn bind(&self) -> Result<(), KrishivError> { + let mut state = self.lock(); + match *state { + State::Unbound => { + *state = State::Waiting; + Ok(()) + } + State::Waiting | State::Done(_) => Err(KrishivObservationReused { + name: self.name.clone(), + } + .into()), + } + } + + /// An action finished: `row` is the metrics of the observed node, if it + /// ran. Only the first action completes the observation. + pub(crate) fn complete(&self, row: Option<&RecordBatch>) { + let mut state = self.lock(); + if matches!(*state, State::Waiting) { + *state = State::Done( + row.cloned() + .unwrap_or_else(|| RecordBatch::new_empty(Arc::new(Schema::empty()))), + ); + self.done.notify_all(); + } + } +} + +/// The metrics of one observed frame (see the module docs). +/// +/// ``` +/// # use krishiv_session::{Observation, Session}; +/// # use krishiv_session::functions::count; +/// # fn main() -> Result<(), krishiv_common::error::KrishivError> { +/// let s = Session::new()?; +/// let obs = Observation::new(); +/// let df = s.range(0, 5, 1, None).observe(&obs, [count(krishiv_session::functions::lit(1)).alias("n")])?; +/// futures::executor::block_on(df.count())?; +/// assert_eq!(obs.get()?.num_rows(), 1); +/// # Ok(()) } +/// ``` +#[derive(Clone)] +pub struct Observation { + pub(crate) shared: Arc, +} + +impl std::fmt::Debug for Observation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Observation") + .field("name", &self.shared.name) + .finish_non_exhaustive() + } +} + +impl Default for Observation { + fn default() -> Self { + Self::new() + } +} + +impl Observation { + /// An observation named by a fresh UUID (Spark's `Observation()`). + #[must_use] + pub fn new() -> Self { + // Unique within the process, which is all a plan's names need. + static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let n = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Self::of(format!("observation-{n}")) + } + + /// An observation named `name` (Spark's `Observation(name)`): the name + /// keys the metrics in the plan, and two relations of one plan may not + /// share it. + /// + /// # Errors + /// `INVALID_OBSERVED_METRICS.MISSING_NAME` for an empty name. + pub fn named(name: &str) -> Result { + if name.is_empty() { + return Err(InvalidObservedMetricsMissingName { + operator: "Observation".to_owned(), + } + .into()); + } + Ok(Self::of(name.to_owned())) + } + + fn of(name: String) -> Self { + Self { + shared: Arc::new(Shared { + name, + state: Mutex::new(State::Unbound), + done: Condvar::new(), + }), + } + } + + /// The observation's name. + #[must_use] + pub fn name(&self) -> &str { + &self.shared.name + } + + /// The metrics, once an action on the observed frame has finished: + /// blocks until then. The first action's metrics stay; later actions do + /// not change them. A frame whose observed node never ran gives a row + /// without columns. + /// + /// # Errors + /// `KRISHIV_NO_OBSERVE_BEFORE_GET` when no frame was observed with it. + pub fn get(&self) -> Result { + self.wait(None)? + .ok_or_else(|| unreachable_wait(&self.shared.name)) + } + + /// As [`Self::get`], giving up after `timeout`: `Ok(None)` if no action + /// has finished by then. + /// + /// # Errors + /// As [`Self::get`]. + pub fn get_timeout(&self, timeout: Duration) -> Result, KrishivError> { + self.wait(Some(timeout)) + } + + fn wait(&self, timeout: Option) -> Result, KrishivError> { + let mut state = self.shared.lock(); + let deadline = timeout.map(|t| std::time::Instant::now() + t); + loop { + match &*state { + State::Unbound => { + return Err(KrishivNoObserveBeforeGet { + name: self.shared.name.clone(), + } + .into()); + } + State::Done(row) => return Ok(Some(row.clone())), + State::Waiting => {} + } + state = match deadline { + None => self + .shared + .done + .wait(state) + .unwrap_or_else(PoisonError::into_inner), + Some(deadline) => { + let left = deadline.saturating_duration_since(std::time::Instant::now()); + if left.is_zero() { + return Ok(None); + } + self.shared + .done + .wait_timeout(state, left) + .unwrap_or_else(PoisonError::into_inner) + .0 + } + }; + } + } +} + +/// `get` waits without a deadline, so `wait` returns a row or an error. +fn unreachable_wait(name: &str) -> KrishivError { + KrishivNoObserveBeforeGet { + name: name.to_owned(), + } + .into() +} diff --git a/crates/krishiv-session/tests/observe.rs b/crates/krishiv-session/tests/observe.rs new file mode 100644 index 000000000..297ceea37 --- /dev/null +++ b/crates/krishiv-session/tests/observe.rs @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `DataFrame::observe` / `Observation` end to end (P3-T084): the classic +//! cases of `tests/fixtures/spark-4.2.0/observe-goldens.json` that do not +//! depend on row-at-a-time pulling, parallelism exactness, first action +//! wins, `LIMIT 0`, reuse errors and cancellation. +#![cfg(test)] + +use std::sync::Arc; +use std::time::Duration; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float64Type, Int32Type, Int64Type}; +use arrow_array::{Array, Float64Array, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::StreamExt; +use futures::executor::block_on; +use krishiv_session::functions::{col, count, lit, max, sum}; +use krishiv_session::{Column, DataFrame, Observation, Session}; +use serde_json::Value; + +/// The generator's `base()`: `k string, v int, x double`. +fn base(s: &Session) -> DataFrame { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Utf8, true), + Field::new("v", DataType::Int32, true), + Field::new("x", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])), + Arc::new(Int32Array::from(vec![Some(1), Some(2), None])), + Arc::new(Float64Array::from(vec![Some(1.5), None, Some(3.0)])), + ], + ) + .unwrap(); + s.register_batches("t", schema, vec![batch]).unwrap(); + s.table("t") +} + +/// One metric cell as the fixture writes it: an integer or null. +fn cell(row: &RecordBatch, i: usize) -> Value { + let c = row.column(i); + if c.is_null(0) { + return Value::Null; + } + match c.data_type() { + DataType::Int64 => Value::from(c.as_primitive::().value(0)), + DataType::Int32 => Value::from(c.as_primitive::().value(0)), + DataType::Float64 => Value::from(c.as_primitive::().value(0)), + DataType::Utf8 | DataType::Utf8View => Value::from(c.as_string_opt::().map_or_else( + || c.as_string_view().value(0).to_owned(), + |s| s.value(0).to_owned(), + )), + other => panic!("unexpected metric type {other}"), + } +} + +fn metrics(row: &RecordBatch) -> serde_json::Map { + row.schema() + .fields() + .iter() + .enumerate() + .map(|(i, f)| (f.name().clone(), cell(row, i))) + .collect() +} + +fn golden(case: &str) -> serde_json::Map { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../tests/fixtures/spark-4.2.0/observe-goldens.json" + ); + let all: Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap(); + all["classic"][case]["metrics"].as_object().unwrap().clone() +} + +/// Runs `df.count()` and returns what the observation holds. +fn run(df: &DataFrame, obs: &Observation) -> serde_json::Map { + block_on(df.count()).unwrap(); + metrics(&obs.get().unwrap()) +} + +/// A golden case: its name, its metrics and what the frame is filtered by below the observation. +type Case = ( + &'static str, + Vec, + Box DataFrame>, +); + +#[test] +fn classic_goldens() { + let s = Session::new().unwrap(); + let t = base(&s); + let n = || count(lit(1)).alias("n"); + let cases: Vec = vec![ + ( + "basic_alias", + vec![ + count(lit(1)).alias("rows"), + sum(col("v")).alias("sv"), + max(col("k")).alias("mk"), + ], + Box::new(|d| d), + ), + ( + "literal_only", + vec![lit(42).alias("c"), lit("s").alias("s")], + Box::new(|d| d), + ), + ( + "filter_before_observe", + vec![n()], + Box::new(|d| d.filter(col("v").eq(1))), + ), + ( + "empty_input", + vec![ + count(lit(1)).alias("c"), + sum(col("v")).alias("s"), + max(col("v")).alias("mx"), + ], + Box::new(|d| d.filter(lit(1).eq(0))), + ), + ( + "nulls_all", + vec![ + sum(col("v")).alias("s"), + count(col("v")).alias("c"), + max(col("v")).alias("m"), + ], + Box::new(|d| d.filter(col("v").is_null())), + ), + ]; + for (name, metrics_of, below) in cases { + let obs = Observation::new(); + let df = below(t.clone()).observe(&obs, metrics_of).unwrap(); + let got = run(&df, &obs); + let want = golden(name); + // `empty_input` and `nulls_all` fixtures list more columns than the + // cases above build; every column built here must agree. + for (k, v) in &got { + assert_eq!(Some(v), want.get(k), "{name}: {k}"); + } + } +} + +#[test] +fn filter_above_counts_the_rows_before_it() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s) + .observe(&obs, [count(lit(1)).alias("n")]) + .unwrap() + .filter(col("v").eq(1)); + assert_eq!(block_on(df.count()).unwrap(), 1); + assert_eq!(metrics(&obs.get().unwrap())["n"], 3); +} + +#[test] +fn unaliased_metrics_are_named_as_spark_names_them() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s) + .observe(&obs, [count(lit(1)), max(col("v")), lit(42)]) + .unwrap(); + let got = run(&df, &obs); + assert_eq!(got["count(1)"], 3); + assert_eq!(got["max(v)"], 2); + assert_eq!(got["42"], 42); +} + +#[test] +fn exact_under_parallelism() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = s + .range(0, 100_000, 1, Some(8)) + .observe( + &obs, + [ + count(lit(1)).alias("n"), + sum(col("id")).alias("s"), + max(col("id")).alias("mx"), + ], + ) + .unwrap(); + let got = run(&df, &obs); + assert_eq!(got["n"], 100_000); + assert_eq!(got["s"], 4_999_950_000_i64); + assert_eq!(got["mx"], 99_999); +} + +#[test] +fn the_first_action_wins() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(df.count()).unwrap(); + let first = metrics(&obs.get().unwrap()); + block_on(df.filter(col("v").eq(1)).count()).unwrap(); + block_on(df.count()).unwrap(); + assert_eq!(metrics(&obs.get().unwrap()), first); + assert_eq!(first["n"], 3); +} + +#[test] +fn limit_zero_reports_nothing() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s) + .observe(&obs, [count(lit(1)).alias("n")]) + .unwrap() + .limit(0); + block_on(df.count()).unwrap(); + assert_eq!(obs.get().unwrap().num_columns(), 0); +} + +#[test] +fn reuse_and_get_before_observe_fail() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + assert_eq!( + obs.get().unwrap_err().error_class().name(), + "KRISHIV_NO_OBSERVE_BEFORE_GET" + ); + let t = base(&s); + t.observe(&obs, [count(lit(1))]).unwrap(); + assert_eq!( + t.observe(&obs, [count(lit(1))]) + .unwrap_err() + .error_class() + .name(), + "KRISHIV_OBSERVATION_REUSED" + ); + // No action yet: `get` waits, and a bounded wait gives up. + assert!( + obs.get_timeout(Duration::from_millis(50)) + .unwrap() + .is_none() + ); +} + +#[test] +fn a_cancelled_action_completes_the_observation() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = s + .range(0, 50_000_000, 1, Some(4)) + .observe(&obs, [count(lit(1)).alias("n")]) + .unwrap(); + block_on(async { + let mut stream = df.execute_stream().await.unwrap(); + let _ = stream.next().await; + drop(stream); + }); + let got = obs + .get_timeout(Duration::from_secs(30)) + .unwrap() + .expect("completes after the stream is dropped"); + // Whatever flowed before the cancellation was counted, no more than the input. + let n = metrics(&got).get("n").and_then(Value::as_i64).unwrap_or(0); + assert!((0..=50_000_000).contains(&n)); +} + +#[test] +fn invalid_metrics_fail_the_action_with_spark_classes() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [col("v")]).unwrap(); + let e = block_on(df.count()).unwrap_err(); + assert_eq!( + e.error_class().name(), + "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE" + ); +} From d7bd15cf6a4a67699c3334536c10439110a9a00e Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 02:45:32 +0000 Subject: [PATCH 05/15] python, connect: DataFrame.observe, Observation and Connect ObservedMetrics (P3-T084 units 5-6) Python: krishiv._native.Observation, DataFrame.observe and krishiv.sql.Observation mirroring pyspark 4.2 (get is a blocking property returning a dict keyed by metric column names; REUSE_OBSERVATION, NO_OBSERVE_BEFORE_GET, VALUE_NOT_NON_EMPTY_STR, NOT_EXPECTED_TYPE, CANNOT_BE_EMPTY). Connect: CollectMetrics translates to PlanNode::CollectMetrics; the server sends ObservedMetrics (name, values, keys, plan_id) after the metrics response of a successful action, none for a failed one. Session gets ObservedSink for handle-less delivery. Coverage inventory updated. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-connect-server/src/execute.rs | 23 ++- .../src/execute/observed.rs | 176 ++++++++++++++++++ .../krishiv-connect-server/tests/execute.rs | 85 +++++++++ crates/krishiv-connect/src/coverage.rs | 2 +- .../krishiv-connect/src/translate/relation.rs | 5 + .../src/translate/relation/tests.rs | 28 +++ crates/krishiv-python/src/dataframe.rs | 51 +++++ crates/krishiv-python/src/lib.rs | 1 + crates/krishiv-session/src/dataframe.rs | 1 + .../krishiv-session/src/dataframe/observe.rs | 81 +++++++- .../krishiv-session/src/dataframe/traced.rs | 2 +- crates/krishiv-session/src/lib.rs | 2 +- docs/reference/connect/coverage.md | 4 +- docs/reference/python/index.md | 1 + .../reference/python/krishiv.sql.dataframe.md | 10 + .../python/krishiv.sql.observation.md | 48 +++++ python/krishiv/_native.pyi | 9 + python/krishiv/errors.py | 5 + python/krishiv/sql/__init__.py | 2 + python/krishiv/sql/dataframe.py | 22 +++ python/krishiv/sql/observation.py | 90 +++++++++ python/tests/test_observe.py | 92 +++++++++ python/tests/test_stubs.py | 1 + tests/connect/coverage.json | 2 +- 24 files changed, 730 insertions(+), 13 deletions(-) create mode 100644 crates/krishiv-connect-server/src/execute/observed.rs create mode 100644 docs/reference/python/krishiv.sql.observation.md create mode 100644 python/krishiv/sql/observation.py create mode 100644 python/tests/test_observe.py diff --git a/crates/krishiv-connect-server/src/execute.rs b/crates/krishiv-connect-server/src/execute.rs index b8be38de1..3ae467e91 100644 --- a/crates/krishiv-connect-server/src/execute.rs +++ b/crates/krishiv-connect-server/src/execute.rs @@ -28,8 +28,10 @@ //! only the envelope is Connect's. pub mod arrow; +pub mod observed; pub mod schema; +use std::collections::BTreeMap; use std::io::Read; use std::sync::Arc; @@ -187,7 +189,7 @@ pub async fn execute_plan( let conf = session.conf(); let empty = pb::Plan::default(); let plan = request.plan.as_ref().unwrap_or(&empty); - let (df, schema) = until_cancelled(&op, async { + let (df, schema, plan_ids) = until_cancelled(&op, async { let translate = TranslateConf { default_data_source: conf .get("spark.sql.sources.default") @@ -214,7 +216,7 @@ pub async fn execute_plan( let logical = translate_plan(plan, decompressed.as_deref(), &translate)?; let df = session.from_plan_with(&logical, QueryOptions::new().operation_id(id.as_str()))?; let schema = df.schema().await?; - Ok((df, schema)) + Ok((df, schema, observed::plan_ids(logical.root()))) }) .await?; let schema_response = schema::proto_schema(&schema)?; @@ -242,6 +244,7 @@ pub async fn execute_plan( schema_response, settings, registered, + plan_ids, )); Ok(futures::stream::unfold(rx, |mut rx| async move { let item = rx.recv().await?; @@ -497,6 +500,7 @@ async fn encode( schema_response: pb::DataType, settings: EncodeSettings, registered: Registered, + plan_ids: BTreeMap, ) { let op = Arc::clone(&sender.op); // The client going away is seen at once, also while the encoder waits @@ -505,7 +509,7 @@ async fn encode( biased; () = op.cancelled() => Err(Stop::Failed(OperationCanceled {}.into())), () = sender.tx.closed() => Err(Stop::Gone), - r = produce(&sender, &df, &schema, schema_response, &settings) => r, + r = produce(&sender, &df, &schema, schema_response, &settings, &plan_ids) => r, }; // The query's stream is gone by now (dropped with `produce`). drop(df); @@ -522,7 +526,10 @@ async fn produce( schema: &SchemaRef, schema_response: pb::DataType, settings: &EncodeSettings, + plan_ids: &BTreeMap, ) -> Result<(), Stop> { + let sink = krishiv_session::ObservedSink::new(); + let df = &df.with_observed_sink(&sink); let mut first = sender.response(None); first.schema = Some(schema_response); sender.send(first).await?; @@ -547,6 +554,16 @@ async fn produce( // metrics (a follow-up); until then the list is empty. metrics.metrics = Some(Metrics::default()); sender.send(metrics).await?; + // The observed nodes' metrics follow the SQL metrics; a query that ended + // in an error never gets here, so it sends none. + for (name, row) in sink.take().unwrap_or_default() { + let plan_id = plan_ids.get(&name).copied().unwrap_or(0); + let mut response = sender.response(None); + response + .observed_metrics + .push(observed::observed_metrics(&name, plan_id, &row)?); + sender.send(response).await?; + } if sender.op.reattachable() { sender .send(sender.response(Some(ResponseType::ResultComplete(ResultComplete {})))) diff --git a/crates/krishiv-connect-server/src/execute/observed.rs b/crates/krishiv-connect-server/src/execute/observed.rs new file mode 100644 index 000000000..c4d9e2117 --- /dev/null +++ b/crates/krishiv-connect-server/src/execute/observed.rs @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `ExecutePlanResponse.observed_metrics` (P3-T084): the metrics of the +//! `CollectMetrics` nodes of a plan, one `ObservedMetrics` per node name +//! (Spark's `SparkConnectPlanExecution.createObservedMetricsResponse`, +//! v4.3.0): the values as literals, the keys the metric column names, the +//! `plan_id` of the relation that carries the node. + +use std::collections::BTreeMap; + +use arrow_array::cast::AsArray as _; +use arrow_array::types::{ + Date32Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, + TimestampMicrosecondType, +}; +use arrow_array::{Array, RecordBatch}; +use arrow_schema::{DataType as Arrow, Field, TimeUnit}; +use krishiv_common::error::{KrishivError, KrishivUnsupportedFeature}; +use krishiv_connect::proto::v4_3::execute_plan_response::ObservedMetrics; +use krishiv_connect::proto::v4_3::expression::Literal; +use krishiv_connect::proto::v4_3::expression::literal::{ + Array as ArrayLiteral, Decimal, LiteralType, +}; +use krishiv_plan::{LogicalPlan, PlanNode}; +use krishiv_types::DataType; + +use super::schema::proto_type; + +/// The plan id of the relation of each observed name in `plan` (the first +/// carrying one; `0` when the relation has none). +pub(crate) fn plan_ids(plan: &LogicalPlan) -> BTreeMap { + let mut out = BTreeMap::new(); + let mut stack = vec![plan]; + while let Some(p) = stack.pop() { + if let PlanNode::CollectMetrics { name, .. } = &p.node { + out.entry(name.clone()) + .or_insert_with(|| p.origin.plan_id.map_or(0, |id| id.0)); + } + p.for_each_child(&mut |c| stack.push(c)); + } + out +} + +/// One `ObservedMetrics` for the one-row `row` of the node `name`; a node that +/// did not run (`LIMIT 0` above it) has no columns and so no values. +/// +/// # Errors +/// A metric of a type that has no literal yet. +pub(crate) fn observed_metrics( + name: &str, + plan_id: i64, + row: &RecordBatch, +) -> Result { + let mut keys = Vec::new(); + let mut values = Vec::new(); + for (field, column) in row.schema().fields().iter().zip(row.columns()) { + keys.push(field.name().clone()); + values.push(literal(field, column.as_ref(), 0)?); + } + Ok(ObservedMetrics { + name: name.to_owned(), + values, + keys, + plan_id, + root_error_idx: None, + errors: Vec::new(), + }) +} + +fn lit(t: LiteralType) -> Literal { + Literal { + literal_type: Some(t), + data_type: None, + } +} + +fn unsupported(t: &Arrow) -> KrishivError { + KrishivUnsupportedFeature { + feature: format!("an observed metric of type {t}"), + } + .into() +} + +fn null(field: &Field) -> Result { + Ok(lit(LiteralType::Null(proto_type( + &DataType::from_arrow_field(field)?, + )?))) +} + +fn literal(field: &Field, a: &dyn Array, i: usize) -> Result { + if a.is_null(i) { + return null(field); + } + Ok(lit(match a.data_type() { + Arrow::Boolean => LiteralType::Boolean(a.as_boolean().value(i)), + Arrow::Int8 => LiteralType::Byte(i32::from(a.as_primitive::().value(i))), + Arrow::Int16 => LiteralType::Short(i32::from(a.as_primitive::().value(i))), + Arrow::Int32 => LiteralType::Integer(a.as_primitive::().value(i)), + Arrow::Int64 => LiteralType::Long(a.as_primitive::().value(i)), + Arrow::Float32 => LiteralType::Float(a.as_primitive::().value(i)), + Arrow::Float64 => LiteralType::Double(a.as_primitive::().value(i)), + Arrow::Utf8 => LiteralType::String(a.as_string::().value(i).to_owned()), + Arrow::LargeUtf8 => LiteralType::String(a.as_string::().value(i).to_owned()), + Arrow::Binary => LiteralType::Binary(a.as_binary::().value(i).to_vec()), + Arrow::LargeBinary => LiteralType::Binary(a.as_binary::().value(i).to_vec()), + Arrow::Date32 => LiteralType::Date(a.as_primitive::().value(i)), + Arrow::Timestamp(TimeUnit::Microsecond, tz) => { + let v = a.as_primitive::().value(i); + if tz.is_some() { + LiteralType::Timestamp(v) + } else { + LiteralType::TimestampNtz(v) + } + } + Arrow::Decimal128(precision, scale) => LiteralType::Decimal(Decimal { + value: a + .as_primitive::() + .value_as_string(i), + precision: Some(i32::from(*precision)), + scale: Some(i32::from(*scale)), + }), + Arrow::List(item) => { + let list = a.as_list::().value(i); + let elements = (0..list.len()) + .map(|j| literal(item, list.as_ref(), j)) + .collect::, _>>()?; + #[expect(deprecated, reason = "Spark's clients still read element_type")] + let array = ArrayLiteral { + element_type: Some(proto_type(&DataType::from_arrow_field(item)?)?), + elements, + }; + LiteralType::Array(array) + } + other => return Err(unsupported(other)), + })) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{Int64Array, StringArray}; + use arrow_schema::Schema; + + use super::*; + + #[test] + fn a_row_becomes_keys_and_literals() { + let schema = Arc::new(Schema::new(vec![ + Field::new("n", Arrow::Int64, false), + Field::new("s", Arrow::Utf8, true), + ])); + let row = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![3])), + Arc::new(StringArray::from(vec![None::<&str>])), + ], + ) + .unwrap(); + let m = observed_metrics("obs", 7, &row).unwrap(); + assert_eq!((m.name.as_str(), m.plan_id), ("obs", 7)); + assert_eq!(m.keys, ["n", "s"]); + assert_eq!(m.values[0].literal_type, Some(LiteralType::Long(3))); + assert!(matches!( + m.values[1].literal_type, + Some(LiteralType::Null(_)) + )); + } + + #[test] + fn a_node_that_did_not_run_has_no_values() { + let row = RecordBatch::new_empty(Arc::new(Schema::empty())); + let m = observed_metrics("obs", 0, &row).unwrap(); + assert!(m.keys.is_empty() && m.values.is_empty()); + } +} diff --git a/crates/krishiv-connect-server/tests/execute.rs b/crates/krishiv-connect-server/tests/execute.rs index 6802ba0dd..11d73578e 100644 --- a/crates/krishiv-connect-server/tests/execute.rs +++ b/crates/krishiv-connect-server/tests/execute.rs @@ -1013,3 +1013,88 @@ async fn large_row_result() { .expect("utf8"); assert_eq!(s.value(0).len(), 65_536); } + +fn literal_long(v: i64) -> pb::Expression { + pb::Expression { + common: None, + expr_type: Some(pb::expression::ExprType::Literal(pb::expression::Literal { + literal_type: Some(pb::expression::literal::LiteralType::Long(v)), + data_type: None, + })), + } +} + +fn aliased_count(alias: &str) -> pb::Expression { + let call = pb::Expression { + common: None, + expr_type: Some(pb::expression::ExprType::UnresolvedFunction( + pb::expression::UnresolvedFunction { + function_name: "count".to_owned(), + arguments: vec![literal_long(1)], + ..Default::default() + }, + )), + }; + pb::Expression { + common: None, + expr_type: Some(pb::expression::ExprType::Alias(Box::new( + pb::expression::Alias { + expr: Some(Box::new(call)), + name: vec![alias.to_owned()], + metadata: None, + }, + ))), + } +} + +fn observed(input: pb::Relation, plan_id: i64) -> pb::Relation { + pb::Relation { + common: Some(pb::RelationCommon { + plan_id: Some(plan_id), + ..Default::default() + }), + rel_type: Some(RelType::CollectMetrics(Box::new(pb::CollectMetrics { + input: Some(Box::new(input)), + name: "m".to_owned(), + metrics: vec![aliased_count("n")], + }))), + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn observed_metrics_follow_the_result() { + let run = start().await; + let responses = collect(&run.channel, request(root(observed(range(10), 7)))) + .await + .expect("execute"); + let metrics: Vec<_> = responses.iter().flat_map(|r| &r.observed_metrics).collect(); + assert_eq!(metrics.len(), 1, "{responses:?}"); + let m = metrics[0]; + assert_eq!((m.name.as_str(), m.plan_id), ("m", 7)); + assert_eq!(m.keys, ["n"]); + assert_eq!( + m.values[0].literal_type, + Some(pb::expression::literal::LiteralType::Long(10)) + ); + // After the arrow batches and the metrics response. + let at = + |f: fn(&pb::ExecutePlanResponse) -> bool| responses.iter().position(f).expect("present"); + assert!(at(|r| !r.observed_metrics.is_empty()) > at(|r| r.metrics.is_some())); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_failed_query_sends_no_observed_metrics() { + let run = start().await; + let failing = pb::Relation { + common: None, + rel_type: Some(RelType::CollectMetrics(Box::new(pb::CollectMetrics { + input: Some(Box::new(sql("SELECT * FROM no_such_table"))), + name: "m".to_owned(), + metrics: vec![aliased_count("n")], + }))), + }; + let status = collect(&run.channel, request(root(failing))) + .await + .expect_err("the table is missing"); + assert_eq!(status.code(), Code::Internal); +} diff --git a/crates/krishiv-connect/src/coverage.rs b/crates/krishiv-connect/src/coverage.rs index 3a2168f50..d8e2662dd 100644 --- a/crates/krishiv-connect/src/coverage.rs +++ b/crates/krishiv-connect/src/coverage.rs @@ -123,7 +123,7 @@ oneof! { relation, RELATION_ROWS, relation::RelType, "Relation.rel_type"; ToSchema => ("to_schema", 26, Missing, "", "lowering: Project of by-name Cast and reorder (Spark's Dataset.to)", "P3-T012", 4), RepartitionByExpression => ("repartition_by_expression", 27, Missing, "", "Repartition (hash or range)", "P3-T012", 3), MapPartitions => ("map_partitions", 28, Missing, "", "new MapPartitions (Python, T014)", "P3-T014", 6), - CollectMetrics => ("collect_metrics", 29, Missing, "", "new CollectMetrics (T084)", "P3-T084", 6), + CollectMetrics => ("collect_metrics", 29, Done, "", "CollectMetrics (observe); the server returns ObservedMetrics after the action", "P3-T084", 6), Parse => ("parse", 30, Missing, "", "lowering: Project of from_json / from_csv over a string relation", "P3-T012", 4), GroupMap => ("group_map", 31, Missing, "", "new GroupMap (applyInPandas, T014)", "P3-T014", 6), CoGroupMap => ("co_group_map", 32, Missing, "", "new CoGroupMap (T014)", "P3-T014", 6), diff --git a/crates/krishiv-connect/src/translate/relation.rs b/crates/krishiv-connect/src/translate/relation.rs index f00ce4ede..c048ad4a0 100644 --- a/crates/krishiv-connect/src/translate/relation.rs +++ b/crates/krishiv-connect/src/translate/relation.rs @@ -95,6 +95,11 @@ impl Translator<'_> { RelType::WithColumns(w) => self.with_columns(w)?, RelType::Range(r) => self.range(r)?, RelType::SetOp(s) => self.set_op(s, origin_of(rel))?, + RelType::CollectMetrics(c) => PlanNode::CollectMetrics { + input: self.input("CollectMetrics", c.input.as_deref())?, + name: c.name.clone(), + metrics: self.expressions(&c.metrics)?, + }, other => { return Err(unsupported_variant( "spark.connect.Relation", diff --git a/crates/krishiv-connect/src/translate/relation/tests.rs b/crates/krishiv-connect/src/translate/relation/tests.rs index 1900d553e..fbed3bc0d 100644 --- a/crates/krishiv-connect/src/translate/relation/tests.rs +++ b/crates/krishiv-connect/src/translate/relation/tests.rs @@ -890,3 +890,31 @@ fn secrets_in_read_options_reach_no_printer_key_or_error() { ); let _ = conf(); } + +#[test] +fn collect_metrics_is_the_observe_node() { + let r = with_id( + relation(RelType::CollectMetrics(Box::new(pb::CollectMetrics { + input: input(), + name: "obs".into(), + metrics: vec![func("count", vec![lit(LiteralType::Integer(1))])], + }))), + 9, + ); + let plan = root(&r); + let PlanNode::CollectMetrics { name, metrics, .. } = &plan.node else { + panic!("a CollectMetrics node, got {:?}", plan.node); + }; + assert_eq!((name.as_str(), metrics.len()), ("obs", 1)); + assert_eq!(plan.origin.plan_id.map(|p| p.0), Some(9)); +} + +#[test] +fn collect_metrics_without_an_input_is_refused() { + let r = relation(RelType::CollectMetrics(Box::new(pb::CollectMetrics { + input: None, + name: "obs".into(), + metrics: vec![], + }))); + fails(translate(&r)); +} diff --git a/crates/krishiv-python/src/dataframe.rs b/crates/krishiv-python/src/dataframe.rs index c78f66ed8..8a31de2b4 100644 --- a/crates/krishiv-python/src/dataframe.rs +++ b/crates/krishiv-python/src/dataframe.rs @@ -26,6 +26,43 @@ pub struct PyDataFrame { pub(crate) owned: Vec>, } +/// `krishiv._native.Observation`: the handle of `DataFrame.observe` (P3-T084). +#[pyclass(name = "Observation", module = "krishiv._native", frozen)] +#[derive(Debug)] +pub struct PyObservation { + inner: krishiv_session::Observation, +} + +#[pymethods] +impl PyObservation { + /// An observation named `name`, or by a fresh unique name. + #[new] + #[pyo3(signature = (name=None))] + fn new(py: Python<'_>, name: Option<&str>) -> PyResult { + let inner = match name { + Some(n) => krishiv_session::Observation::named(n).map_err(|e| to_py(py, &e))?, + None => krishiv_session::Observation::new(), + }; + Ok(Self { inner }) + } + + /// The name keying the metrics. + #[getter] + fn name(&self) -> String { + self.inner.name().to_owned() + } + + /// The one-row metrics, as a stream: blocks (the GIL released) until + /// the first action on the observed frame has finished. + fn get(&self, py: Python<'_>) -> PyResult { + let obs = self.inner.clone(); + let batch = py.detach(|| obs.get()).map_err(|e| to_py(py, &e))?; + let schema = batch.schema(); + let reader = arrow_array::RecordBatchIterator::new([Ok(batch)], schema); + Ok(ArrowStream::new(Box::new(reader))) + } +} + /// A temporary table dropped from its session when the last holder goes. #[derive(Debug)] pub(crate) struct OwnedTable { @@ -171,6 +208,20 @@ impl PyDataFrame { self.derive(self.inner.order_by(columns(cols))) } + /// Observes `metrics` under `observation` (PySpark's `observe`). + fn observe( + &self, + py: Python<'_>, + observation: &PyObservation, + metrics: Vec, + ) -> PyResult { + let inner = self + .inner + .observe(&observation.inner, columns(metrics)) + .map_err(|e| to_py(py, &e))?; + Ok(self.derive(inner)) + } + fn limit(&self, n: i32) -> Self { self.derive(self.inner.limit(n)) } diff --git a/crates/krishiv-python/src/lib.rs b/crates/krishiv-python/src/lib.rs index 9d8023c57..bb16d5a6f 100644 --- a/crates/krishiv-python/src/lib.rs +++ b/crates/krishiv-python/src/lib.rs @@ -38,6 +38,7 @@ fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_function(wrap_pyfunction!(column::column, m)?)?; diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index 321899562..b4775694d 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -32,6 +32,7 @@ pub(crate) mod checkpoint; mod fill_stream; mod geo_gate; mod observe; +pub use observe::ObservedSink; mod traced; pub(crate) use fill_stream::{CacheFill, FillKind}; diff --git a/crates/krishiv-session/src/dataframe/observe.rs b/crates/krishiv-session/src/dataframe/observe.rs index 24f605250..4ba06b4eb 100644 --- a/crates/krishiv-session/src/dataframe/observe.rs +++ b/crates/krishiv-session/src/dataframe/observe.rs @@ -3,8 +3,9 @@ //! its [`Observation`] handles when an action's result stream ends (P3-T084). //! A child of `dataframe` so it reads the frame's private fields. +use std::collections::BTreeMap; use std::pin::Pin; -use std::sync::Arc; +use std::sync::{Arc, Mutex, PoisonError}; use std::task::{Context, Poll}; use arrow_array::RecordBatch; @@ -59,6 +60,24 @@ impl DataFrame { }) } + /// This frame, reporting the metrics of its observed nodes to `sink` when + /// an action's result has been read to its end (Spark Connect's + /// `ObservedMetrics`; the server has no [`Observation`] handles). + #[must_use] + pub fn with_observed_sink(&self, sink: &ObservedSink) -> Self { + let mut out = self.clone(); + out.keep.push(Arc::clone(&sink.slot) as _); + out + } + + /// The sinks this frame carries. + pub(super) fn sinks(&self) -> Vec> { + self.keep + .iter() + .filter_map(|k| Arc::clone(k).downcast::().ok()) + .collect() + } + /// The observation handles this frame carries. pub(super) fn observations(&self) -> Vec> { self.keep @@ -68,6 +87,42 @@ impl DataFrame { } } +/// The state of an [`ObservedSink`]. +pub(crate) type SinkSlot = Mutex>>; + +/// Where an action leaves the metrics of the observed nodes of its plan, by +/// name (one row each): set when the result was read to its end without an +/// error, replaced by each later action (Spark Connect: the last action +/// wins, a failed query sends none). +#[derive(Clone, Default)] +pub struct ObservedSink { + slot: Arc, +} + +impl std::fmt::Debug for ObservedSink { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ObservedSink").finish_non_exhaustive() + } +} + +impl ObservedSink { + /// An empty sink. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// The metrics the last successful action left, once; `None` when no + /// action has finished without an error. + #[must_use] + pub fn take(&self) -> Option> { + self.slot + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + } +} + /// `stream`, completing `observers` from `query`'s observations when it ends, /// fails or is dropped (Spark completes them whether the action succeeded or /// not); the first action wins. @@ -75,14 +130,17 @@ pub(super) fn deliver( stream: RecordBatchStream, query: Arc, observers: Vec>, + sinks: Vec>, ) -> RecordBatchStream { - if observers.is_empty() { + if observers.is_empty() && sinks.is_empty() { return stream; } Box::pin(Delivering { inner: stream, query, observers, + sinks, + finished: false, }) } @@ -90,15 +148,27 @@ struct Delivering { inner: RecordBatchStream, query: Arc, observers: Vec>, + sinks: Vec>, + /// The stream ended without an error. + finished: bool, } impl Delivering { fn complete(&mut self) { let observers = std::mem::take(&mut self.observers); - if observers.is_empty() { + let sinks = std::mem::take(&mut self.sinks); + if observers.is_empty() && sinks.is_empty() { return; } - let observed = self.query.observed().unwrap_or_default(); + let observed = self.query.observed(); + if self.finished + && let Ok(all) = &observed + { + for sink in sinks { + *sink.lock().unwrap_or_else(PoisonError::into_inner) = Some(all.clone()); + } + } + let observed = observed.unwrap_or_default(); for o in observers { o.complete(observed.get(o.name())); } @@ -110,6 +180,9 @@ impl Stream for Delivering { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let polled = self.inner.as_mut().poll_next(cx); + if matches!(polled, Poll::Ready(None)) { + self.finished = true; + } if matches!(polled, Poll::Ready(None | Some(Err(_)))) { self.complete(); } diff --git a/crates/krishiv-session/src/dataframe/traced.rs b/crates/krishiv-session/src/dataframe/traced.rs index ab8e3c70f..4492448f6 100644 --- a/crates/krishiv-session/src/dataframe/traced.rs +++ b/crates/krishiv-session/src/dataframe/traced.rs @@ -74,7 +74,7 @@ impl DataFrame { .run_with_query(&physical, &a.env) .instrument(exec.clone()) .await?; - let stream = super::observe::deliver(stream, query, self.observations()); + let stream = super::observe::deliver(stream, query, self.observations(), self.sinks()); let stream = crate::telemetry::hold_span(stream, exec); let stream = fill.fill(stream, &self.session.state.cache); let stream = match result { diff --git a/crates/krishiv-session/src/lib.rs b/crates/krishiv-session/src/lib.rs index 2c880ec0f..a7186fbac 100644 --- a/crates/krishiv-session/src/lib.rs +++ b/crates/krishiv-session/src/lib.rs @@ -110,7 +110,7 @@ pub use column::Column; pub use conf::RuntimeConf; pub use dataframe::{ DataFrame, DataFrameReader, DataFrameWriter, DataFrameWriterV2, ExplainMode, GroupedData, - JoinType, StorageLevel, + JoinType, ObservedSink, StorageLevel, }; pub use explain::Estimate; pub use explain_graph::{ diff --git a/docs/reference/connect/coverage.md b/docs/reference/connect/coverage.md index f8e702aba..db4bab3f9 100644 --- a/docs/reference/connect/coverage.md +++ b/docs/reference/connect/coverage.md @@ -5,7 +5,7 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve ## Relation.rel_type -62 arms: 7 done, 6 partial, 45 missing, 4 out of scope. +62 arms: 8 done, 6 partial, 44 missing, 4 out of scope. | # | Arm | Status | Refused fields | IR target | Task | Group | |---:|---|---|---|---|---|---:| @@ -36,7 +36,7 @@ Every arm of the `oneof` fields the Connect translator dispatches on, for the ve | 26 | to_schema | missing | | lowering: Project of by-name Cast and reorder (Spark's Dataset.to) | P3-T012 | 4 | | 27 | repartition_by_expression | missing | | Repartition (hash or range) | P3-T012 | 3 | | 28 | map_partitions | missing | | new MapPartitions (Python, T014) | P3-T014 | 6 | -| 29 | collect_metrics | missing | | new CollectMetrics (T084) | P3-T084 | 6 | +| 29 | collect_metrics | done | | CollectMetrics (observe); the server returns ObservedMetrics after the action | P3-T084 | 6 | | 30 | parse | missing | | lowering: Project of from_json / from_csv over a string relation | P3-T012 | 4 | | 31 | group_map | missing | | new GroupMap (applyInPandas, T014) | P3-T014 | 6 | | 32 | co_group_map | missing | | new CoGroupMap (T014) | P3-T014 | 6 | diff --git a/docs/reference/python/index.md b/docs/reference/python/index.md index 03b671abf..eff2363bf 100644 --- a/docs/reference/python/index.md +++ b/docs/reference/python/index.md @@ -13,6 +13,7 @@ | [`krishiv.sql.functions`](krishiv.sql.functions.md) | `pyspark.sql.functions`: column builders. A function resolves by name in the engine's | | [`krishiv.sql.group`](krishiv.sql.group.md) | `GroupedData`: a DataFrame grouped by key columns, as `pyspark.sql.GroupedData`. | | [`krishiv.sql.na`](krishiv.sql.na.md) | `DataFrameNaFunctions` (`DataFrame.na`): handling of missing data, as `pyspark.sql.DataFrameNaFunctions`. | +| [`krishiv.sql.observation`](krishiv.sql.observation.md) | `Observation`: named metrics of a DataFrame, collected as an action runs (`pyspark.sql.Observation`, 4.2). | | [`krishiv.sql.pandas`](krishiv.sql.pandas.md) | `pyspark.sql.pandas`: the pandas and Arrow vectorized UDFs (`krishiv.sql.pandas.functions`). | | [`krishiv.sql.pandas.functions`](krishiv.sql.pandas.functions.md) | `pyspark.sql.pandas.functions`: `pandas_udf` and `arrow_udf`. | | [`krishiv.sql.profile_view`](krishiv.sql.profile_view.md) | An offline viewer for `krishiv.query-profile` documents. | diff --git a/docs/reference/python/krishiv.sql.dataframe.md b/docs/reference/python/krishiv.sql.dataframe.md index 79706033f..484280d18 100644 --- a/docs/reference/python/krishiv.sql.dataframe.md +++ b/docs/reference/python/krishiv.sql.dataframe.md @@ -825,6 +825,16 @@ def na(self) -> DataFrameNaFunctions Functionality for working with missing data (`DataFrameNaFunctions`). +### `observe` + +```python +def observe(self, observation: Observation | str, *exprs: Column) -> DataFrame +``` + +Observes named metrics (aggregate expressions) of this DataFrame as the first action runs +(`DataFrame.observe`). With an `krishiv.sql.Observation` they are read from +`observation.get`; with a name only, the metrics are computed with the query. + ### `offset` ```python diff --git a/docs/reference/python/krishiv.sql.observation.md b/docs/reference/python/krishiv.sql.observation.md new file mode 100644 index 000000000..71e646de0 --- /dev/null +++ b/docs/reference/python/krishiv.sql.observation.md @@ -0,0 +1,48 @@ +# `krishiv.sql.observation` + + + +`Observation`: named metrics of a DataFrame, collected as an action runs (`pyspark.sql.Observation`, 4.2). + +The metrics are aggregate expressions given to `DataFrame.observe`; they are computed over the rows that +flow through that point of the plan while the first action runs, and `Observation.get` returns them as a +`dict` keyed by the metric column names. + +## `Observation` + +```python +class Observation +``` + +Metrics of one DataFrame, collected as its first action runs. + +One `Observation` observes one DataFrame. `name` is optional (a unique name is made up) and only keys the +metrics inside the plan. + +```python +>>> from krishiv.sql import Observation, functions as F +>>> observation = Observation("stats") # doctest: +SKIP +>>> df.observe(observation, F.count(F.lit(1)).alias("rows")).count() # doctest: +SKIP +>>> observation.get # doctest: +SKIP +{'rows': 100} +``` + +### `get` + +```python +@property +def get(self) -> dict[str, Any] +``` + +The metrics as a `dict` keyed by the metric column names, once the first action on the observed +DataFrame has finished (blocks until then, as Spark; later actions do not change it). An observed plan +that never ran its metrics (`limit(0)` above it) gives `{}`. + +### `name` + +```python +@property +def name(self) -> str +``` + +The name keying the metrics in the plan (not in PySpark's public API; kept for the Connect translation). diff --git a/python/krishiv/_native.pyi b/python/krishiv/_native.pyi index ece448ffa..92e0a587f 100644 --- a/python/krishiv/_native.pyi +++ b/python/krishiv/_native.pyi @@ -20,6 +20,7 @@ __all__ = [ "Column", "DataFrame", "GroupedData", + "Observation", "Session", "WindowSpec", "__allocator__", @@ -66,6 +67,13 @@ class Session: def refresh_files(self) -> None: ... def cache_reserved(self) -> int: ... +@final +class Observation: + def __init__(self, name: str | None = ...) -> None: ... + @property + def name(self) -> str: ... + def get(self) -> ArrowStream: ... + @final class DataFrame: def select(self, cols: list[Column]) -> DataFrame: ... @@ -79,6 +87,7 @@ class DataFrame: self, other: DataFrame, on: Column | None = ..., using: list[str] | None = ..., how: str = ... ) -> DataFrame: ... def order_by(self, cols: list[Column]) -> DataFrame: ... + def observe(self, observation: Observation, metrics: list[Column]) -> DataFrame: ... def limit(self, n: int) -> DataFrame: ... def alias(self, alias: str) -> DataFrame: ... def col(self, name: str) -> Column: ... diff --git a/python/krishiv/errors.py b/python/krishiv/errors.py index 214af8ee0..63cdf502a 100644 --- a/python/krishiv/errors.py +++ b/python/krishiv/errors.py @@ -55,6 +55,11 @@ "NOT_NUMERIC_COLUMNS": ( "Numeric aggregation function can only be applied on numeric columns, got ." ), + "NOT_LIST_OF_COLUMN": "Argument `` should be a list[Column], got .", + "REUSE_OBSERVATION": "An Observation can be used with a DataFrame only once.", + "NO_OBSERVE_BEFORE_GET": "Should observe by calling `DataFrame.observe` before `get`.", + "VALUE_NOT_NON_EMPTY_STR": "Value for `` must be a non-empty string, got ''.", + "CANNOT_BE_EMPTY": "At least one must be specified.", "VALUE_NOT_POSITIVE": "Value for `` must be positive, got ''.", "NEGATIVE_VALUE": "Value for `` must be greater than or equal to 0, got ''.", "ARGUMENT_REQUIRED": "Argument `` is required when .", diff --git a/python/krishiv/sql/__init__.py b/python/krishiv/sql/__init__.py index 62e188dbd..c97fd1ee0 100644 --- a/python/krishiv/sql/__init__.py +++ b/python/krishiv/sql/__init__.py @@ -7,6 +7,7 @@ from krishiv.sql.dataframe import DataFrame from krishiv.sql.group import GroupedData from krishiv.sql.na import DataFrameNaFunctions +from krishiv.sql.observation import Observation from krishiv.sql.readwriter import DataFrameReader, DataFrameWriter, DataFrameWriterV2 from krishiv.sql.row import Row from krishiv.sql.session import Catalog, SparkSession @@ -24,6 +25,7 @@ "DataFrameWriter", "DataFrameWriterV2", "GroupedData", + "Observation", "Row", "RuntimeConfig", "SparkSession", diff --git a/python/krishiv/sql/dataframe.py b/python/krishiv/sql/dataframe.py index fcd791543..1b77a3e61 100644 --- a/python/krishiv/sql/dataframe.py +++ b/python/krishiv/sql/dataframe.py @@ -17,6 +17,7 @@ from krishiv.errors import ( AnalysisException, PySparkAttributeError, + PySparkException, PySparkTypeError, PySparkValueError, not_expected_type, @@ -27,6 +28,7 @@ from krishiv.sql.functions import first as _first from krishiv.sql.group import GroupedData from krishiv.sql.na import _NO_VALUE, DataFrameNaFunctions +from krishiv.sql.observation import Observation, _non_empty, _pyspark_error from krishiv.sql.profile_view import ProfileView from krishiv.sql.row import Row from krishiv.sql.stat import DataFrameStatFunctions @@ -426,6 +428,26 @@ def orderBy(self, *cols: ColumnsOrNames, ascending: bool | list[bool] = True) -> sort = orderBy + def observe(self, observation: Observation | str, *exprs: Column) -> DataFrame: + """Observes named metrics (aggregate expressions) of this DataFrame as the first action runs + (``DataFrame.observe``). With an :class:`~krishiv.sql.Observation` they are read from + ``observation.get``; with a name only, the metrics are computed with the query.""" + if not isinstance(observation, Observation | str): + raise not_expected_type("observation", "Observation or str", observation) + if len(exprs) == 0: + raise PySparkValueError(errorClass="CANNOT_BE_EMPTY", messageParameters={"item": "exprs"}) + if not all(isinstance(c, Column) for c in exprs): + raise PySparkTypeError( + errorClass="NOT_LIST_OF_COLUMN", messageParameters={"arg_name": "exprs", "arg_type": "tuple"} + ) + if isinstance(observation, str): + _non_empty(observation) + observation = Observation(observation) + try: + return self._new(self._df.observe(observation._native, [c._c for c in exprs])) + except PySparkException as e: + raise _pyspark_error(e) from e + def limit(self, num: int) -> DataFrame: if not isinstance(num, int): raise not_expected_type("num", "int", num) diff --git a/python/krishiv/sql/observation.py b/python/krishiv/sql/observation.py new file mode 100644 index 000000000..50b80a630 --- /dev/null +++ b/python/krishiv/sql/observation.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``Observation``: named metrics of a DataFrame, collected as an action runs (``pyspark.sql.Observation``, 4.2). + +The metrics are aggregate expressions given to :meth:`DataFrame.observe`; they are computed over the rows that +flow through that point of the plan while the first action runs, and :attr:`Observation.get` returns them as a +``dict`` keyed by the metric column names. +""" + +from __future__ import annotations + +from typing import Any + +from krishiv import _native +from krishiv.errors import ( + PySparkAssertionError, + PySparkException, + PySparkValueError, + not_expected_type, +) +from krishiv.sql import _convert + +__all__ = ["Observation"] + +# The engine's Krishiv-only classes and the PySpark ones the client raises for them. +_REUSED = "KRISHIV_OBSERVATION_REUSED" +_NO_OBSERVE = "KRISHIV_NO_OBSERVE_BEFORE_GET" + + +def _pyspark_error(error: PySparkException) -> PySparkException: + """``error`` as PySpark's client raises it: the Krishiv-only classes become the PySpark ones.""" + condition = error.getCondition() + if condition == _REUSED: + return PySparkAssertionError(errorClass="REUSE_OBSERVATION", messageParameters={}) + if condition == _NO_OBSERVE: + return PySparkAssertionError(errorClass="NO_OBSERVE_BEFORE_GET", messageParameters={}) + return error + + +def _non_empty(name: str) -> None: + if name == "": + raise PySparkValueError( + errorClass="VALUE_NOT_NON_EMPTY_STR", messageParameters={"arg_name": "name", "arg_value": name} + ) + + +class Observation: + """Metrics of one DataFrame, collected as its first action runs. + + One ``Observation`` observes one DataFrame. ``name`` is optional (a unique name is made up) and only keys the + metrics inside the plan. + + >>> from krishiv.sql import Observation, functions as F + >>> observation = Observation("stats") # doctest: +SKIP + >>> df.observe(observation, F.count(F.lit(1)).alias("rows")).count() # doctest: +SKIP + >>> observation.get # doctest: +SKIP + {'rows': 100} + """ + + def __init__(self, name: str | None = None) -> None: + if name is not None: + if not isinstance(name, str): + raise not_expected_type("name", "str", name) + _non_empty(name) + try: + self._native = _native.Observation(name) + except PySparkException as e: # an empty name only; checked above + raise _pyspark_error(e) from e + + @property + def name(self) -> str: + """The name keying the metrics in the plan (not in PySpark's public API; kept for the Connect translation).""" + return self._native.name + + @property + def get(self) -> dict[str, Any]: + """The metrics as a ``dict`` keyed by the metric column names, once the first action on the observed + DataFrame has finished (blocks until then, as Spark; later actions do not change it). An observed plan + that never ran its metrics (``limit(0)`` above it) gives ``{}``.""" + try: + stream = self._native.get() + except PySparkException as e: + raise _pyspark_error(e) from e + pa = _convert.require_pyarrow("Observation.get") + table = pa.table(stream) + if table.num_columns == 0 or table.num_rows == 0: + return {} + return _convert.rows_of(table)[0].asDict() + + def __repr__(self) -> str: + return f"Observation({self.name!r})" diff --git a/python/tests/test_observe.py b/python/tests/test_observe.py new file mode 100644 index 000000000..e440737f9 --- /dev/null +++ b/python/tests/test_observe.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``DataFrame.observe`` and ``Observation`` (P3-T084), against the Spark 4.2 behaviour.""" + +from __future__ import annotations + +import threading + +import pytest + +from krishiv.errors import PySparkAssertionError, PySparkTypeError, PySparkValueError +from krishiv.sql import DataFrame, Observation, SparkSession +from krishiv.sql import functions as F + +pytest.importorskip("pyarrow") + + +@pytest.fixture +def df() -> DataFrame: + return SparkSession.builder.create().range(0, 10) + + +def test_basic(df: DataFrame) -> None: + obs = Observation("stats") + out = df.observe(obs, F.count(F.lit(1)).alias("n"), F.max("id").alias("m")) + assert out.count() == 10 + assert obs.get == {"n": 10, "m": 9} + + +def test_unnamed_and_unaliased_names(df: DataFrame) -> None: + obs = Observation() + df.observe(obs, F.count(F.lit(1)), F.max("id")).collect() + assert obs.get == {"count(1)": 10, "max(id)": 9} + + +def test_first_action_wins(df: DataFrame) -> None: + obs = Observation() + out = df.observe(obs, F.count(F.lit(1)).alias("n")) + out.count() + out.filter(F.col("id") > 5).count() + assert obs.get == {"n": 10} + + +def test_get_blocks_until_action(df: DataFrame) -> None: + obs = Observation() + out = df.observe(obs, F.count(F.lit(1)).alias("n")) + seen = [] + t = threading.Thread(target=lambda: seen.append(obs.get)) + t.start() + t.join(0.2) + assert t.is_alive() + out.count() + t.join(10) + assert seen == [{"n": 10}] + + +def test_string_name_form(df: DataFrame) -> None: + assert df.observe("m", F.count(F.lit(1)).alias("n")).count() == 10 + + +def test_reuse(df: DataFrame) -> None: + obs = Observation() + df.observe(obs, F.count(F.lit(1))) + with pytest.raises(PySparkAssertionError) as e: + df.observe(obs, F.count(F.lit(1))) + assert e.value.getCondition() == "REUSE_OBSERVATION" + + +def test_get_before_observe() -> None: + with pytest.raises(PySparkAssertionError) as e: + _ = Observation().get + assert e.value.getCondition() == "NO_OBSERVE_BEFORE_GET" + + +def test_argument_errors(df: DataFrame) -> None: + with pytest.raises(PySparkValueError) as e: + Observation("") + assert e.value.getCondition() == "VALUE_NOT_NON_EMPTY_STR" + with pytest.raises(PySparkTypeError) as t: + Observation(5) # ty: ignore[invalid-argument-type] + assert t.value.getCondition() == "NOT_EXPECTED_TYPE" + with pytest.raises(PySparkValueError) as c: + df.observe(Observation()) + assert c.value.getCondition() == "CANNOT_BE_EMPTY" + with pytest.raises(PySparkValueError) as c2: + df.observe("m") + assert c2.value.getCondition() == "CANNOT_BE_EMPTY" + + +def test_limit_zero_gives_empty(df: DataFrame) -> None: + obs = Observation() + df.observe(obs, F.count(F.lit(1)).alias("n")).limit(0).collect() + assert obs.get == {} diff --git a/python/tests/test_stubs.py b/python/tests/test_stubs.py index 8fb75a47b..bf873bccf 100644 --- a/python/tests/test_stubs.py +++ b/python/tests/test_stubs.py @@ -87,6 +87,7 @@ def test_native_stub_matches_runtime() -> None: "WindowSpec", "ArrowStream", "AnonymousPythonUdf", + "Observation", } for cls_name, methods in classes.items(): runtime = getattr(krishiv_native, cls_name) diff --git a/tests/connect/coverage.json b/tests/connect/coverage.json index a48c98101..b324d1de7 100644 --- a/tests/connect/coverage.json +++ b/tests/connect/coverage.json @@ -26,7 +26,7 @@ {"oneof": "Relation.rel_type", "field": "to_schema", "number": 26, "status": "missing", "refused": "", "target": "lowering: Project of by-name Cast and reorder (Spark's Dataset.to)", "task": "P3-T012", "group": 4}, {"oneof": "Relation.rel_type", "field": "repartition_by_expression", "number": 27, "status": "missing", "refused": "", "target": "Repartition (hash or range)", "task": "P3-T012", "group": 3}, {"oneof": "Relation.rel_type", "field": "map_partitions", "number": 28, "status": "missing", "refused": "", "target": "new MapPartitions (Python, T014)", "task": "P3-T014", "group": 6}, - {"oneof": "Relation.rel_type", "field": "collect_metrics", "number": 29, "status": "missing", "refused": "", "target": "new CollectMetrics (T084)", "task": "P3-T084", "group": 6}, + {"oneof": "Relation.rel_type", "field": "collect_metrics", "number": 29, "status": "done", "refused": "", "target": "CollectMetrics (observe); the server returns ObservedMetrics after the action", "task": "P3-T084", "group": 6}, {"oneof": "Relation.rel_type", "field": "parse", "number": 30, "status": "missing", "refused": "", "target": "lowering: Project of from_json / from_csv over a string relation", "task": "P3-T012", "group": 4}, {"oneof": "Relation.rel_type", "field": "group_map", "number": 31, "status": "missing", "refused": "", "target": "new GroupMap (applyInPandas, T014)", "task": "P3-T014", "group": 6}, {"oneof": "Relation.rel_type", "field": "co_group_map", "number": 32, "status": "missing", "refused": "", "target": "new CoGroupMap (T014)", "task": "P3-T014", "group": 6}, From 625b4632b2f9251f9fe512373d777e26e9050825 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 05:45:57 +0000 Subject: [PATCH 06/15] observe: every action completes the observation; plan-stable instance, limit cap, pool-charged state (P3-T084 review fixes) One delivery hook (run_on_query_logged) completes observations and Connect sinks for reads, writes, table writes, collect_with_metrics and profiles; a plan that observes is not served from the result cache. CollectMetrics carries a planner-assigned instance (codec field 4) instead of a process counter, and a limit (field 5) set from a LIMIT above so counts are rows taken, as Spark's. The operator charges its state to the driver's consumer. A scalar subquery in a metric is KRISHIV_UNSUPPORTED_FEATURE. A failed finalization is what Observation.get returns. Connect sends all names in one response. public-api snapshots updated; internal items are doc(hidden). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-connect-server/src/execute.rs | 12 ++- crates/krishiv-exec/CLAUDE.md | 1 + crates/krishiv-exec/src/aggregate/observe.rs | 95 ++++++++++++++++-- crates/krishiv-exec/src/context.rs | 38 ++++--- .../krishiv-exec/src/physical/instantiate.rs | 17 +++- crates/krishiv-exec/src/physical/mod.rs | 12 ++- crates/krishiv-exec/src/planner/build.rs | 10 +- crates/krishiv-exec/src/planner/observe.rs | 20 ++++ crates/krishiv-plan-codec/src/decode.rs | 2 + crates/krishiv-plan-codec/src/encode.rs | 4 + .../src/proto/krishiv.plan.v1.pb.rs | 7 ++ .../tests/common/inventory.rs | 2 + .../tests/golden/observe/observe.bin | Bin 893 -> 897 bytes crates/krishiv-plan/src/analyzer/observe.rs | 27 ++++- crates/krishiv-session/CLAUDE.md | 1 + crates/krishiv-session/public-api.txt | 23 ++++- crates/krishiv-session/src/dataframe.rs | 2 + .../krishiv-session/src/dataframe/observe.rs | 80 +++++++++------ .../krishiv-session/src/dataframe/traced.rs | 5 +- crates/krishiv-session/src/observation.rs | 29 +++--- crates/krishiv-session/src/result_cache.rs | 9 ++ crates/krishiv-session/src/session.rs | 12 ++- crates/krishiv-session/src/session/env.rs | 3 + .../python/krishiv.sql.observation.md | 3 +- proto/krishiv/plan/v1/plan.proto | 5 + python/krishiv/sql/dataframe.py | 11 +- python/krishiv/sql/observation.py | 3 +- python/public-api.txt | 18 +++- 28 files changed, 347 insertions(+), 104 deletions(-) diff --git a/crates/krishiv-connect-server/src/execute.rs b/crates/krishiv-connect-server/src/execute.rs index 3ae467e91..3364ca8ce 100644 --- a/crates/krishiv-connect-server/src/execute.rs +++ b/crates/krishiv-connect-server/src/execute.rs @@ -554,14 +554,16 @@ async fn produce( // metrics (a follow-up); until then the list is empty. metrics.metrics = Some(Metrics::default()); sender.send(metrics).await?; - // The observed nodes' metrics follow the SQL metrics; a query that ended - // in an error never gets here, so it sends none. + // The observed nodes' metrics follow the SQL metrics, all in one response as Spark sends + // them; a query that ended in an error never gets here, so it sends none. + let mut all = Vec::new(); for (name, row) in sink.take().unwrap_or_default() { let plan_id = plan_ids.get(&name).copied().unwrap_or(0); + all.push(observed::observed_metrics(&name, plan_id, &row)?); + } + if !all.is_empty() { let mut response = sender.response(None); - response - .observed_metrics - .push(observed::observed_metrics(&name, plan_id, &row)?); + response.observed_metrics = all; sender.send(response).await?; } if sender.op.reattachable() { diff --git a/crates/krishiv-exec/CLAUDE.md b/crates/krishiv-exec/CLAUDE.md index b0527c390..a3f85cf4a 100644 --- a/crates/krishiv-exec/CLAUDE.md +++ b/crates/krishiv-exec/CLAUDE.md @@ -70,6 +70,7 @@ failed, never a short result), `stream.rs` (a partition evaluated from disk: lea memory), `disk.rs` (frames and `lag`/`lead`/`nth_value` of a partition on disk: forward readers, a frame's batches held or folded from the file), `group_limit.rs` (Spark's `WindowGroupLimit`: the exact limit operator below the window, the `Ranker` the sort also uses to prune each sorted batch while that pays; the planner infers it for a filter over the window, `planner/window.rs`), `tests.rs` (property test against a naive reference; spilled vs in-memory equality). `RANGE` value bounds are `FrameBound::RangeInt`/`RangeFloat` (the key's units), `RangeDayTime` (a day-time interval over a `TIMESTAMP`/`TIMESTAMP_NTZ` key: whole days on the zone's clock, then elapsed time) and `RangeMonths` (calendar months over a `DATE` or `TIMESTAMP` key; the `TIMESTAMP` zone is the session's, passed to `Window::new` because the column carries only a marker); `month_tests.rs` (month bounds from disk vs memory in a DST zone). +`aggregate/observe.rs` (P3-T084) `Dataset.observe`: `PhysNode::CollectMetrics { input, name, measure, instance, limit }`, a pass-through operator that folds each driver's rows into a global partial `AggTable` (charged to the driver's consumer; `NeedMemory` under pressure) and forwards the batch unchanged; `limit` (set by `planner/observe.rs` `cap_observed` from a `LIMIT` above, through projections only) caps the rows folded, one budget per instance, as Spark counts the rows its consumer pulled. Drivers' states go to `QueryContext::observe_partial` keyed `(name, instance)` (the planner's plan-order rank, plan data; of one name the highest wins); `QueryContext::observed()` merges and finalizes them. `LIMIT 0` plans no node. Operators with work after their input (`needs_flush`) get `end_input` + a per-pipeline barrier + `flush`. `loops.rs` (T180 unit 5b) `UnionLoopSource`: a recursive CTE (`PhysNode::UnionLoop`/`LoopRef`) as one source. The anchor and each round of the recursive term run as queries of their own over a child `QueryContext` (child memory pool, cancelled with the diff --git a/crates/krishiv-exec/src/aggregate/observe.rs b/crates/krishiv-exec/src/aggregate/observe.rs index c570d3275..9debe8cef 100644 --- a/crates/krishiv-exec/src/aggregate/observe.rs +++ b/crates/krishiv-exec/src/aggregate/observe.rs @@ -15,11 +15,22 @@ //! //! Two instances of one name in a plan (a self-join of an observed //! relation) do not add up: the last one's states are the observation's -//! (Spark's `CollectMetricsExec.collect` keeps the last). +//! (Spark's `CollectMetricsExec.collect` keeps the last). The rank of an +//! instance is plan data (`PhysNode::CollectMetrics::instance`, set by the +//! planner), so workers that each instantiate a part of the plan agree on it. +//! +//! A `LIMIT` above (through row-for-row operators) caps what is folded at the +//! rows it takes (`limit`, one budget for the instance's drivers): Spark +//! counts the rows its consumer pulled, not the batches a pull read. The +//! states are charged to the query's memory pool through the driver's +//! consumer, so a `collect_list` metric over a large input fails with +//! `UNABLE_TO_ACQUIRE_MEMORY` instead of growing without bound. use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use krishiv_memory::Reservation; + use arrow_array::RecordBatch; use arrow_schema::SchemaRef; use krishiv_types::batch::Batch; @@ -32,16 +43,12 @@ use crate::hash::HashSeed; use crate::operator::{LocalState, Operator, OperatorOutput}; use crate::operators::{Project, local_mut}; -/// Orders the instances of a plan (instantiation is one thread, in plan -/// order): the highest is the last in the plan. -static INSTANCES: AtomicU64 = AtomicU64::new(0); - /// Turns the partial states of one observation into its row. #[derive(Debug)] pub(crate) struct Finalizer { name: String, /// This instance's place among the plan's. - instance: u64, + instance: u32, /// Merges the states of the drivers and gives the aggregates' results. core: Arc, /// The metrics over those results, one column each. @@ -55,7 +62,7 @@ impl Finalizer { } /// The instance's rank in its plan. - pub(crate) fn instance(&self) -> u64 { + pub(crate) fn instance(&self) -> u32 { self.instance } @@ -84,6 +91,10 @@ pub(crate) struct CollectMetrics { /// The partial aggregate (no keys). core: Arc, finalizer: Arc, + /// The rows a `LIMIT` above takes, if one does. + limit: Option, + /// Rows folded so far, over the instance's drivers (counts against `limit`). + taken: AtomicU64, } impl CollectMetrics { @@ -91,8 +102,14 @@ impl CollectMetrics { /// arguments (`pre_schema`: its output, or `schema` without it), /// `partial` is the global aggregate's plan, `post` computes the /// metrics from its results. + #[expect( + clippy::too_many_arguments, + reason = "the operator's plan data and builders" + )] pub(crate) fn try_new( name: &str, + instance: u32, + limit: Option, schema: SchemaRef, pre: Option>, partial: super::AggregateSpec, @@ -109,7 +126,7 @@ impl CollectMetrics { let post = post(&merge.output_schema)?; let finalizer = Arc::new(Finalizer { name: name.to_owned(), - instance: INSTANCES.fetch_add(1, Ordering::Relaxed), + instance, core: merge, post, }); @@ -118,6 +135,8 @@ impl CollectMetrics { pre, core, finalizer, + limit, + taken: AtomicU64::new(0), })) } } @@ -129,6 +148,11 @@ struct Local { finalizer: Arc, partition: u32, sent: bool, + /// The table's bytes, charged to the driver's consumer. + reservation: Reservation, + /// The rows of the input batch being folded, when a retry after `NeedMemory` must not + /// claim them from the budget again. + claimed: Option, } impl Local { @@ -153,6 +177,26 @@ impl Drop for Local { } } +impl CollectMetrics { + /// How many of `rows` just read count: all, or what is left of the limit above's budget. + fn claim(&self, rows: usize) -> usize { + let Some(cap) = self.limit else { + return rows; + }; + let mut got = 0; + // `try_update` retries the closure on contention, so `got` is the last attempt's. + let _ = self + .taken + .try_update(Ordering::Relaxed, Ordering::Relaxed, |taken| { + got = usize::try_from(cap.saturating_sub(taken)) + .unwrap_or(usize::MAX) + .min(rows); + Some(taken.saturating_add(got as u64)) + }); + got + } +} + impl Operator for CollectMetrics { fn lane_safe(&self) -> bool { true @@ -173,6 +217,8 @@ impl Operator for CollectMetrics { finalizer: Arc::clone(&self.finalizer), partition: ctx.partition(), sent: false, + reservation: ctx.consumer().reservation(), + claimed: None, })) } @@ -181,16 +227,45 @@ impl Operator for CollectMetrics { if input.num_rows() == 0 { return Ok(OperatorOutput::NeedMoreInput(None)); } + let take = match l.claimed { + Some(take) => take, + None => self.claim(input.num_rows()), + }; + if take == 0 { + // A limit above has all the rows it takes: these are not counted. + return Ok(OperatorOutput::NeedMoreInput(Some(input.clone()))); + } + l.claimed = Some(take); + let counted = if take < input.num_rows() { + Some(input.slice(0, take)?) + } else { + None + }; + let seen = counted.as_ref().unwrap_or(input); let args; let folded = match &self.pre { Some(pre) => { - args = pre.project(input, l.partition)?; + args = pre.project(seen, l.partition)?; &args } - None => input, + None => seen, }; let p = l.table.prepare(folded, false)?; + let want = l.table.projected(&p); + let need = want.saturating_sub(l.reservation.size()); + if need > 0 && l.reservation.try_grow(need).is_err() { + return Ok(OperatorOutput::NeedMemory { bytes: need }); + } l.table.apply(&p)?; + l.claimed = None; + // A fold that outgrew its projection is charged now. + let keep = l.table.memory_size(); + let over = keep.saturating_sub(l.reservation.size()); + if over > 0 { + l.reservation.try_grow(over)?; + } + l.reservation + .shrink(l.reservation.size().saturating_sub(keep)); Ok(OperatorOutput::NeedMoreInput(Some(input.clone()))) } diff --git a/crates/krishiv-exec/src/context.rs b/crates/krishiv-exec/src/context.rs index 48f43fe37..835f872b0 100644 --- a/crates/krishiv-exec/src/context.rs +++ b/crates/krishiv-exec/src/context.rs @@ -226,9 +226,9 @@ pub struct QueryContext { /// The scans' decode working sets (`krishiv_io::ScanGate`), made on /// first use with a budget of a quarter of the query's memory limit. scan_gate: OnceLock>, - /// The partial states of each observation (`Dataset.observe`), by name - /// and instance (see [`Self::observe_partial`]). - observed: Mutex>>, + /// The partial states of each observing node (`Dataset.observe`), by name + /// and the planner's instance number (see [`Self::observe_partial`]). + observed: Mutex>, finished: AtomicBool, finished_event: Event, /// Blocking closures in flight; the query ends only once they have. @@ -292,18 +292,14 @@ impl QueryContext { finalizer: &Arc, states: arrow_array::RecordBatch, ) { - let mut all = lock(&self.observed); - let instances = all.entry(finalizer.name().to_owned()).or_default(); - match instances - .iter_mut() - .find(|o| o.finalizer.instance() == finalizer.instance()) - { - Some(o) => o.partials.push(states), - None => instances.push(Observed { + lock(&self.observed) + .entry((finalizer.name().to_owned(), finalizer.instance())) + .or_insert_with(|| Observed { finalizer: Arc::clone(finalizer), - partials: vec![states], - }), - } + partials: Vec::new(), + }) + .partials + .push(states); } /// The observations (`Dataset.observe`) of the query so far: each name's @@ -315,15 +311,15 @@ impl QueryContext { /// /// # Errors /// A metric's own error (finalizing runs its expressions). + #[doc(hidden)] pub fn observed(&self) -> Result> { let all = lock(&self.observed); - let mut out = BTreeMap::new(); - for (name, instances) in all.iter() { - if let Some(last) = instances.iter().max_by_key(|o| o.finalizer.instance()) { - out.insert(name.clone(), last.finalizer.finalize(&last.partials)?); - } - } - Ok(out) + // Keys sort by name then instance: the last one of a name overwrites. + let last: BTreeMap<&String, &Observed> = + all.iter().map(|((name, _), o)| (name, o)).collect(); + last.into_iter() + .map(|(name, o)| Ok((name.clone(), o.finalizer.finalize(&o.partials)?))) + .collect() } /// The query id. diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index f1bc9918e..1c6e7b7b6 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -1039,8 +1039,10 @@ impl Instantiation<'_> { input, name, measure, + instance, + limit, } => self.with_subqueries(input, &[], |input, v| { - collect_metrics_op(name, measure, input, v) + collect_metrics_op(name, *instance, *limit, measure, input, v) }), PhysNode::ReplicateRows { input, @@ -1704,6 +1706,8 @@ pub(crate) fn expand_op( /// shape of `measure` is what `planner/observe.rs` makes. pub(crate) fn collect_metrics_op( name: &str, + instance: u32, + limit: Option, measure: &PhysNode, schema: &SchemaRef, v: &Values<'_>, @@ -1730,8 +1734,13 @@ pub(crate) fn collect_metrics_op( else { return Err(bad()); }; - if !subqueries.is_empty() - || !plan.group_by.is_empty() + if !subqueries.is_empty() { + return Err(krishiv_common::error::KrishivUnsupportedFeature { + feature: "a subquery in an observed metric (Dataset.observe)".to_owned(), + } + .into()); + } + if !plan.group_by.is_empty() || plan.output != AggregateOutput::Final || plan.distinct.iter().any(|d| *d) { @@ -1759,6 +1768,8 @@ pub(crate) fn collect_metrics_op( let merge = final_of(&partial)?; CollectMetrics::try_new( name, + instance, + limit, Arc::clone(schema), pre, aggregate_spec(&pre_schema, &partial, v.env, None), diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index 584e6abf6..dcce6aeed 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -436,6 +436,15 @@ pub enum PhysNode { name: String, /// The metrics' computation over the input's columns. measure: Box, + /// The planner's rank of this node among the plan's observing + /// nodes, in plan order (Spark's `CollectMetricsExec.collect` lets + /// the last of a name win). Plan data, so every worker numbers the + /// instances alike. + instance: u32, + /// At most this many rows are folded: a `LIMIT` above (through + /// row-for-row operators only) takes no more, and Spark counts the + /// rows its consumer pulled. + limit: Option, }, /// Each row of `input` `count` times (`INTERSECT ALL`, `EXCEPT ALL`). ReplicateRows { @@ -1037,7 +1046,7 @@ plan_data!(enum PhysNode { Filter { input, layout, predicate, subqueries }, Project { input, layout, items, nullable, subqueries }, Limit { input, offset, limit }, - CollectMetrics { input, name, measure }, + CollectMetrics { input, name, measure, instance, limit }, ReplicateRows { input, layout, count, subqueries }, PythonEval { input, layout, calls }, Generate { input, layout, kind, outer, arg, subqueries }, @@ -1161,6 +1170,7 @@ impl PhysNode { /// that is not a child (it reads no relation), listed for the codec and /// the feature scan. #[must_use] + #[doc(hidden)] pub fn measure(&self) -> Option<&PhysNode> { match self { Self::CollectMetrics { measure, .. } => Some(measure), diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index 272252594..29bcdc1f4 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -312,6 +312,9 @@ pub(crate) struct Planner<'a> { /// Nesting depth of the relations whose rows are never read (under a /// `LIMIT 0`): `observe` plans no node there. pub(super) unobserved: u32, + /// The rank the next observing node takes (`PhysNode::CollectMetrics::instance`), in plan + /// order. + pub(super) observers: u32, } impl<'a> Planner<'a> { @@ -347,6 +350,7 @@ impl<'a> Planner<'a> { loops: HashMap::new(), moves: super::build_project::Moves::default(), unobserved: 0, + observers: 0, } } @@ -796,7 +800,11 @@ impl<'a> Planner<'a> { } return Ok(Self::skip_rows(top, skip, n)); } - let input = self.plan(input)?; + let mut input = self.plan(input)?; + // Spark counts the rows an observed relation was pulled for: the rows a limit above takes. + if let Some(total) = skip.checked_add(n) { + super::observe::cap_observed(&mut input.node, total); + } // Each partition keeps at most `skip + n` rows before the gather; // the skip applies only after it. let input = if input.placement.satisfies(Distribution::Single) { diff --git a/crates/krishiv-exec/src/planner/observe.rs b/crates/krishiv-exec/src/planner/observe.rs index 22a7496f3..085b77fe2 100644 --- a/crates/krishiv-exec/src/planner/observe.rs +++ b/crates/krishiv-exec/src/planner/observe.rs @@ -29,6 +29,9 @@ impl Planner<'_> { name: &str, metrics: &[Arc], ) -> Result { + // Numbered before its input's nodes: Spark's plan order is pre-order. + let instance = self.observers; + self.observers += 1; let planned = self.plan(input)?; if self.unobserved > 0 { return Ok(planned); @@ -52,6 +55,8 @@ impl Planner<'_> { input, name: name.to_owned(), measure: Box::new(measured.node), + instance, + limit: None, }, schema, layout, @@ -62,3 +67,18 @@ impl Planner<'_> { Ok(out) } } + +/// Caps the rows the observing nodes of `node` fold at `rows`, through the operators that pass +/// every row on (a projection, another observing node): a `LIMIT` above takes no more than +/// that, and Spark's `CollectMetricsExec` counts the rows its consumer pulled. A filter, sort, +/// join or any operator that changes which rows arrive ends the search. +pub(super) fn cap_observed(node: &mut PhysNode, rows: u64) { + match node { + PhysNode::Project { input, .. } => cap_observed(input, rows), + PhysNode::CollectMetrics { input, limit, .. } => { + *limit = Some(limit.map_or(rows, |l| l.min(rows))); + cap_observed(input, rows); + } + _ => {} + } +} diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index b3a9e8335..eb7ba1350 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -536,6 +536,8 @@ impl Decoder { input: self.boxed(m.input)?, name: m.name.clone(), measure: self.boxed(m.measure)?, + instance: m.instance, + limit: m.limit, }, Kind::ReplicateRows(r) => PhysNode::ReplicateRows { input: self.boxed(r.input)?, diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 49c77cadb..a71e1d3c8 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -314,10 +314,14 @@ impl Encoder { input, name, measure, + instance, + limit, } => Kind::CollectMetrics(pb::CollectMetricsNode { input: self.node_index(input)?, name: name.clone(), measure: self.node_index(measure)?, + instance: *instance, + limit: *limit, }), PhysNode::ReplicateRows { input, diff --git a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index 5f3cf79d3..3a4185ee5 100644 --- a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs +++ b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs @@ -234,6 +234,13 @@ pub struct CollectMetricsNode { pub name: ::prost::alloc::string::String, #[prost(uint32, tag = "3")] pub measure: u32, + /// The planner's rank of this node among the plan's observing nodes, in plan + /// order: of several nodes with one name the highest is the observation's. + #[prost(uint32, tag = "4")] + pub instance: u32, + /// At most this many rows are folded (a `LIMIT` above takes no more). + #[prost(uint64, optional, tag = "5")] + pub limit: ::core::option::Option, } /// feature: python_udf /// The rows of `input`, then one column per Python UDF call over every row. diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index 852e0b002..d5ac51b2c 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -1700,6 +1700,8 @@ pub(crate) fn collect_metrics() -> PhysNode { PhysNode::CollectMetrics { input: boxed(scan()), name: "obs".to_owned(), + instance: 1, + limit: Some(5), measure: boxed(PhysNode::Project { input: boxed(PhysNode::Aggregate { input: boxed(PhysNode::OneRow), diff --git a/crates/krishiv-plan-codec/tests/golden/observe/observe.bin b/crates/krishiv-plan-codec/tests/golden/observe/observe.bin index cce7548ac3eb9c727e73b123dbf23f430daabbf5..ef495f9ccb977ff3f0f1e91f40fa105ecc4030dd 100644 GIT binary patch delta 36 scmey%*2uo0k5Pp079+P1bAD2>1hWF82CD*#)&!=B3q>|>Vm!?V0Jp;l?EnA( delta 24 fcmZo<|I41oP&(j87N=Ud0D& diff --git a/crates/krishiv-plan/src/analyzer/observe.rs b/crates/krishiv-plan/src/analyzer/observe.rs index 7ab3e641b..a29cb7232 100644 --- a/crates/krishiv-plan/src/analyzer/observe.rs +++ b/crates/krishiv-plan/src/analyzer/observe.rs @@ -79,6 +79,18 @@ fn violation(metric: &Arc) -> Option { None } +/// Whether `e` holds a subquery expression. +fn has_subquery(e: &Arc) -> bool { + let mut stack = vec![Arc::clone(e)]; + while let Some(e) = stack.pop() { + if matches!(e.node, ExprNode::Subquery { .. }) { + return true; + } + e.for_each_child(&mut |k| stack.push(Arc::clone(k))); + } + false +} + impl Resolver<'_> { /// Resolves `metrics` over `input` as a global aggregate would and /// keeps the node with them bound; its output is the input's. @@ -92,9 +104,10 @@ impl Resolver<'_> { ) -> Result { let o = p.origin; let c = self.child(input, env)?; - self.observing = true; + // Restored, not cleared: a metric's own subquery may observe too. + let outer = std::mem::replace(&mut self.observing, true); let items = self.select_items(&c, metrics, env, "CollectMetrics"); - self.observing = false; + self.observing = outer; let items = items?; let zone = self.ctx.conf.session_zone; let shown = |b: &super::exprs::Bound| format!("\"{}\"", pretty(&b.shown(), zone)); @@ -142,6 +155,16 @@ impl Resolver<'_> { }; return Err(AnalysisError::new(e, o)); } + // Spark accepts a scalar subquery in a metric; the engine does not run one beside the + // rows that pass through, and says so rather than failing inside the planner. + if items.iter().any(|b| has_subquery(&b.expr)) { + return Err(AnalysisError::new( + krishiv_common::error::KrishivUnsupportedFeature { + feature: "a subquery in an observed metric (Dataset.observe)".to_owned(), + }, + o, + )); + } // One relation per name: the same relation (its canonical // representative) may appear again, another may not. let key = addr(&*self.rep(p)); diff --git a/crates/krishiv-session/CLAUDE.md b/crates/krishiv-session/CLAUDE.md index fc423d955..7692c435c 100644 --- a/crates/krishiv-session/CLAUDE.md +++ b/crates/krishiv-session/CLAUDE.md @@ -33,6 +33,7 @@ Paths are relative to this crate. - `src/char_write.rs` the write side of `CHAR(n)` / `VARCHAR(n)` columns (trim trailing spaces to the limit, else `EXCEED_LIMIT_LENGTH`; off under `spark.sql.legacy.charVarcharAsString`; inside arrays, maps and structs too) and `output_columns` (CTAS and views keep the declared type); `tests/char_varchar_write.rs` - `src/result_cache.rs` the result cache of an engine runtime (`krishiv.resultCache.*`, off by default; design `design/P2-T094-result-cache.md`; design P2-E1 R6: sessions built alone share the process-wide store `global()`, a supplied runtime owns its own, `EngineRuntime::results`; `krishiv.resultCache.maxBytes` is a runtime key): `CacheKey` plus the identity of every permanent table read, a query over permanent catalog tables with no non-deterministic call is served from and stored into it at `DataFrame::execute_stream`; dropped by the relation cache's invalidation calls (`RelationCache::clear` / `invalidate_table`: REFRESH TABLE, CLEAR CACHE, engine writes) and by the TTL. Temporary views over permanent tables are inlined (eligible); in-memory tables, reads by path and plans that substituted a session-cached relation are not. - `src/dataframe.rs` `DataFrameWriter` (files: `save`, `parquet`/`csv`/`json`/`text`; tables: `save_as_table`, `insert_into`; `bucket_by`/`sort_by`, which only Spark's `saveAsTable` takes and the engine has no bucketed tables for) and `DataFrameWriterV2` (`writeTo`). Table writers have no statement of their own: they run `table_commands.rs` `exec_save_as_table` / `exec_insert_into` / `exec_write_v2` (one path with `CREATE TABLE AS` and `INSERT`) and make the access checks of the statement they stand for through `authorize_equivalent` (`tests/access_policy.rs`). `saveAsTable` overwrite drops and recreates the table only after every check that can refuse the new one (`validate_new_table`, `Session::check_write_ready`, `UNSUPPORTED_OVERWRITE.TABLE`); `tests/dataframe_writer_tables.rs` +- `src/observation.rs` + `src/dataframe/observe.rs` (P3-T084) `Observation` (first action wins, `get` blocks, an error finalizing the metrics is what `get` returns) and `DataFrame::observe`; a frame's observations and Connect `ObservedSink`s travel in `DataFrame::keep` into `QueryEnv::observers`, and `run_on_query_logged` wraps the result stream of every run (reads, writes, metrics, profiles) in `dataframe::deliver`, which completes them when it ends, fails or is dropped. A plan with a `CollectMetrics` is never served from the result cache (`result_cache.rs` `Found::observed`); `tests/observe.rs`, `tests/observe_goldens.rs` (the fixture runner with `KNOWN_GAPS`) - `src/explain.rs` EXPLAIN and `DataFrame::estimate` (ADR 0041: estimated table-scan bytes and rows from the optimizer's statistics, `None` when unknown; the CLI's agent cost gate uses it); `src/profile.rs` profile for EXPLAIN ANALYZE - `src/cache.rs` the session relation cache (`cache`/`persist`/`unpersist`: registrations, fill on use, invalidation; design P3-T083); `cache/tier.rs` where an entry lives (`HeldTable` in memory, its pool bytes held until the last reading query ends; `SpilledTable` on disk in a spill scope of its own, `CacheDisk`) and `CacheSpill`, the `relation_cache` consumer's `Spillable` (least recently used unread entries spill, `MEMORY_AND_DISK`, or are evicted, `MEMORY_ONLY`); an action holds the snapshots it reads in a `SnapshotLease`; `cache/level.rs` `StorageLevel::from_name` and the session default (`spark.sql.defaultCacheStorageLevel`); `cache/checkpoint_stats.rs` what `system.caches` shows of checkpoints; `dataframe/checkpoint.rs` `checkpoint`/`local_checkpoint` (the snapshot is a hidden table owned by the frame, as a script result is; `eager = false` leaves a `LazyCheckpoint` in `keep` that `plan_with_checkpoints` takes at the first analysis; reliable ones also write a spill-format file under `spark.checkpoint.dir`, deleted with the snapshot); `tests/cache.rs`, `tests/checkpoint.rs` - `src/system.rs` the reserved `system` catalog (P2-T126): `system.` with an empty namespace, rows built per query from process state (one split in-process); `system.capabilities`, `system.metrics` (the engine metrics snapshot: labels `MAP`, histogram `buckets` `ARRAY`), `system.caches` (`system/caches.rs`: the file-metadata cache from `MetadataCache::stats()` and the relation cache, with memory and disk bytes, pending registrations, spills and evictions), `system.operators` (one row per operator of a finished query, from its `QueryMetrics`), `system.queries` (the runtime's `QueryHistory` seen through the session's `QueryLog`: UUIDv7 or caller operation id, state, text, times, wall time, peak memory, error class; recorded around every query in `run_on_query`; the principal's own rows unless it holds `MONITOR` on `system`) and `system.memory` (the pool tree from `MemoryPool::snapshot()`, taken when the scan is planned) so far (`register_build_features` is how the binary lists its cargo features); `query_text` follows `krishiv.system.queryText.maxBytes` (default 4096, cut at a character boundary, `query_text_truncated`) and `krishiv.system.queryText.literals` (`keep`, or `mask`: `krishiv_sql::mask_literals`, every literal `?`), read from the session configuration when a statement is recorded (`TextPolicy`); `CREATE|DROP` of anything in `system` is `KRISHIV_RESERVED_NAME` (`access.rs`); `tests/system_tables.rs` diff --git a/crates/krishiv-session/public-api.txt b/crates/krishiv-session/public-api.txt index c5b42981a..c058a4f02 100644 --- a/crates/krishiv-session/public-api.txt +++ b/crates/krishiv-session/public-api.txt @@ -592,7 +592,6 @@ impl krishiv_session::DataFrame pub fn krishiv_session::DataFrame::agg(&self, impl core::iter::traits::collect::IntoIterator) -> Self pub fn krishiv_session::DataFrame::alias(&self, &str) -> Self pub fn krishiv_session::DataFrame::cache(&self) -> Self -pub async fn krishiv_session::DataFrame::checkpoint(&self, bool) -> core::result::Result pub fn krishiv_session::DataFrame::col(&self, &str) -> krishiv_session::Column pub async fn krishiv_session::DataFrame::collect(&self) -> core::result::Result, krishiv_common::error::KrishivError> pub async fn krishiv_session::DataFrame::count(&self) -> core::result::Result @@ -615,7 +614,6 @@ pub async fn krishiv_session::DataFrame::is_cached(&self) -> core::result::Resul pub fn krishiv_session::DataFrame::join(&self, &Self, krishiv_session::Column, krishiv_session::JoinType) -> Self pub fn krishiv_session::DataFrame::join_using<'a>(&self, &Self, impl core::iter::traits::collect::IntoIterator, krishiv_session::JoinType) -> Self pub fn krishiv_session::DataFrame::limit(&self, i32) -> Self -pub async fn krishiv_session::DataFrame::local_checkpoint(&self, bool, core::option::Option) -> core::result::Result pub fn krishiv_session::DataFrame::offset(&self, i32) -> Self pub fn krishiv_session::DataFrame::order_by(&self, impl core::iter::traits::collect::IntoIterator>) -> Self pub fn krishiv_session::DataFrame::persist(&self, krishiv_session::StorageLevel) -> Self @@ -642,6 +640,9 @@ pub fn krishiv_session::DataFrame::write_to_table(&self, &str) -> krishiv_sessio pub async fn krishiv_session::DataFrame::write_to_table_with(&self, &str, core::option::Option, alloc::vec::Vec<(alloc::string::String, alloc::string::String)>, alloc::vec::Vec<(alloc::string::String, alloc::string::String)>, alloc::vec::Vec, alloc::vec::Vec, &str, core::option::Option) -> core::result::Result<(), krishiv_common::error::KrishivError> pub async fn krishiv_session::DataFrame::write_to_with(&self, &str, &str, &str, alloc::vec::Vec<(alloc::string::String, alloc::string::String)>, alloc::vec::Vec, core::option::Option<(i32, alloc::vec::Vec)>, core::option::Option>) -> core::result::Result<(), krishiv_common::error::KrishivError> impl krishiv_session::DataFrame +pub async fn krishiv_session::DataFrame::checkpoint(&self, bool) -> core::result::Result +pub async fn krishiv_session::DataFrame::local_checkpoint(&self, bool, core::option::Option) -> core::result::Result +impl krishiv_session::DataFrame pub async fn krishiv_session::DataFrame::collect_with_metrics(&self) -> core::result::Result<(alloc::vec::Vec, krishiv_exec::metrics::QueryMetrics), krishiv_common::error::KrishivError> impl krishiv_session::DataFrame pub async fn krishiv_session::DataFrame::collect_with_profile(&self) -> core::result::Result<(alloc::vec::Vec, krishiv_exec::profile::QueryProfile), krishiv_common::error::KrishivError> @@ -654,6 +655,8 @@ impl krishiv_session::DataFrame pub async fn krishiv_session::DataFrame::estimate(&self) -> core::result::Result impl krishiv_session::DataFrame pub async fn krishiv_session::DataFrame::explain_with(&self, krishiv_session::ExplainMode, krishiv_session::GraphFormat) -> core::result::Result +impl krishiv_session::DataFrame +pub fn krishiv_session::DataFrame::observe(&self, &krishiv_session::Observation, impl core::iter::traits::collect::IntoIterator>) -> core::result::Result impl core::clone::Clone for krishiv_session::DataFrame pub fn krishiv_session::DataFrame::clone(&self) -> krishiv_session::DataFrame impl core::fmt::Debug for krishiv_session::DataFrame @@ -763,6 +766,19 @@ impl core::clone::Clone for krishiv_session::GroupedData pub fn krishiv_session::GroupedData::clone(&self) -> krishiv_session::GroupedData impl core::fmt::Debug for krishiv_session::GroupedData pub fn krishiv_session::GroupedData::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct krishiv_session::Observation +impl krishiv_session::Observation +pub fn krishiv_session::Observation::get(&self) -> core::result::Result +pub fn krishiv_session::Observation::get_timeout(&self, core::time::Duration) -> core::result::Result, krishiv_common::error::KrishivError> +pub fn krishiv_session::Observation::name(&self) -> &str +pub fn krishiv_session::Observation::named(&str) -> core::result::Result +pub fn krishiv_session::Observation::new() -> Self +impl core::clone::Clone for krishiv_session::Observation +pub fn krishiv_session::Observation::clone(&self) -> krishiv_session::Observation +impl core::default::Default for krishiv_session::Observation +pub fn krishiv_session::Observation::default() -> Self +impl core::fmt::Debug for krishiv_session::Observation +pub fn krishiv_session::Observation::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct krishiv_session::OptimizerStep pub krishiv_session::OptimizerStep::batch: alloc::string::String pub krishiv_session::OptimizerStep::diff: alloc::string::String @@ -999,7 +1015,6 @@ pub krishiv_session::StorageLevel::use_disk: bool pub krishiv_session::StorageLevel::use_memory: bool pub krishiv_session::StorageLevel::use_off_heap: bool impl krishiv_session::StorageLevel -pub fn krishiv_session::StorageLevel::from_name(&str) -> core::option::Option pub const krishiv_session::StorageLevel::DISK_ONLY: Self pub const krishiv_session::StorageLevel::DISK_ONLY_2: Self pub const krishiv_session::StorageLevel::MEMORY_AND_DISK: Self @@ -1008,6 +1023,8 @@ pub const krishiv_session::StorageLevel::MEMORY_ONLY: Self pub const krishiv_session::StorageLevel::MEMORY_ONLY_2: Self pub const krishiv_session::StorageLevel::NONE: Self pub const krishiv_session::StorageLevel::OFF_HEAP: Self +impl krishiv_session::StorageLevel +pub fn krishiv_session::StorageLevel::from_name(&str) -> core::option::Option impl core::clone::Clone for krishiv_session::StorageLevel pub fn krishiv_session::StorageLevel::clone(&self) -> krishiv_session::StorageLevel impl core::cmp::Eq for krishiv_session::StorageLevel diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index b4775694d..092100e74 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -33,6 +33,7 @@ mod fill_stream; mod geo_gate; mod observe; pub use observe::ObservedSink; +pub(crate) use observe::{Observers, deliver}; mod traced; pub(crate) use fill_stream::{CacheFill, FillKind}; @@ -1052,6 +1053,7 @@ impl DataFrame { if let Some(t) = text { env.set_text(t); } + env.observers = Arc::new(self.observers()); let analysis = Arc::new(Analysis { env: Arc::new(env), analyzed: Arc::new(analyzed), diff --git a/crates/krishiv-session/src/dataframe/observe.rs b/crates/krishiv-session/src/dataframe/observe.rs index 4ba06b4eb..79174bdac 100644 --- a/crates/krishiv-session/src/dataframe/observe.rs +++ b/crates/krishiv-session/src/dataframe/observe.rs @@ -44,6 +44,7 @@ impl DataFrame { /// As [`Self::observe`] with a name and no handle (Spark's /// `Dataset.observe(String, ...)`): the metrics are computed with the /// query and reported to whoever reads the query's observations. + #[doc(hidden)] #[must_use] pub fn observe_named( &self, @@ -63,28 +64,37 @@ impl DataFrame { /// This frame, reporting the metrics of its observed nodes to `sink` when /// an action's result has been read to its end (Spark Connect's /// `ObservedMetrics`; the server has no [`Observation`] handles). + /// + /// For the Connect server: not part of the engine's embedding API. + #[doc(hidden)] #[must_use] pub fn with_observed_sink(&self, sink: &ObservedSink) -> Self { let mut out = self.clone(); out.keep.push(Arc::clone(&sink.slot) as _); + // The analysis made so far did not carry the sink. + out.analysis = Arc::default(); out } - /// The sinks this frame carries. - pub(super) fn sinks(&self) -> Vec> { - self.keep - .iter() - .filter_map(|k| Arc::clone(k).downcast::().ok()) - .collect() + /// The observations and sinks this frame carries, for the query's env. + pub(super) fn observers(&self) -> Observers { + let mut out = Observers::default(); + for k in &self.keep { + if let Ok(shared) = Arc::clone(k).downcast::() { + out.observations.push(shared); + } else if let Ok(sink) = Arc::clone(k).downcast::() { + out.sinks.push(sink); + } + } + out } +} - /// The observation handles this frame carries. - pub(super) fn observations(&self) -> Vec> { - self.keep - .iter() - .filter_map(|k| Arc::clone(k).downcast::().ok()) - .collect() - } +/// What a frame's runs report their observed metrics to. +#[derive(Default)] +pub(crate) struct Observers { + observations: Vec>, + sinks: Vec>, } /// The state of an [`ObservedSink`]. @@ -93,7 +103,8 @@ pub(crate) type SinkSlot = Mutex>>; /// Where an action leaves the metrics of the observed nodes of its plan, by /// name (one row each): set when the result was read to its end without an /// error, replaced by each later action (Spark Connect: the last action -/// wins, a failed query sends none). +/// wins, a failed query sends none). For the Connect server. +#[doc(hidden)] #[derive(Clone, Default)] pub struct ObservedSink { slot: Arc, @@ -125,42 +136,45 @@ impl ObservedSink { /// `stream`, completing `observers` from `query`'s observations when it ends, /// fails or is dropped (Spark completes them whether the action succeeded or -/// not); the first action wins. -pub(super) fn deliver( +/// not); the first action wins. The one hook every run of a query passes. +pub(crate) fn deliver( stream: RecordBatchStream, query: Arc, - observers: Vec>, - sinks: Vec>, + observers: &Observers, ) -> RecordBatchStream { - if observers.is_empty() && sinks.is_empty() { + if observers.observations.is_empty() && observers.sinks.is_empty() { return stream; } Box::pin(Delivering { inner: stream, query, - observers, - sinks, + observers: Some((observers.observations.clone(), observers.sinks.clone())), finished: false, }) } +/// The observations and sinks a stream completes. +type Targets = (Vec>, Vec>); + struct Delivering { inner: RecordBatchStream, query: Arc, - observers: Vec>, - sinks: Vec>, + observers: Option, /// The stream ended without an error. finished: bool, } impl Delivering { fn complete(&mut self) { - let observers = std::mem::take(&mut self.observers); - let sinks = std::mem::take(&mut self.sinks); - if observers.is_empty() && sinks.is_empty() { + let Some((observations, sinks)) = self.observers.take() else { return; - } - let observed = self.query.observed(); + }; + // On the carriers a caller reads (decimals, intervals), as the action's rows are. + let observed = self.query.observed().and_then(|all| { + all.into_iter() + .map(|(name, row)| Ok((name, crate::carrier::client_batch(&row)?))) + .collect::, KrishivError>>() + }); if self.finished && let Ok(all) = &observed { @@ -168,9 +182,13 @@ impl Delivering { *sink.lock().unwrap_or_else(PoisonError::into_inner) = Some(all.clone()); } } - let observed = observed.unwrap_or_default(); - for o in observers { - o.complete(observed.get(o.name())); + for o in observations { + o.complete( + observed + .as_ref() + .map(|all| all.get(o.name())) + .map_err(Clone::clone), + ); } } } diff --git a/crates/krishiv-session/src/dataframe/traced.rs b/crates/krishiv-session/src/dataframe/traced.rs index 4492448f6..647b1c215 100644 --- a/crates/krishiv-session/src/dataframe/traced.rs +++ b/crates/krishiv-session/src/dataframe/traced.rs @@ -69,12 +69,11 @@ impl DataFrame { // The execution span stays open while the rows are produced and read: it closes when // the result stream is dropped, not when `run` returns the stream. let exec = tracing::info_span!("krishiv.execute"); - let (stream, query) = self + let stream = self .session - .run_with_query(&physical, &a.env) + .run(&physical, &a.env) .instrument(exec.clone()) .await?; - let stream = super::observe::deliver(stream, query, self.observations(), self.sinks()); let stream = crate::telemetry::hold_span(stream, exec); let stream = fill.fill(stream, &self.session.state.cache); let stream = match result { diff --git a/crates/krishiv-session/src/observation.rs b/crates/krishiv-session/src/observation.rs index 7bdc52daa..ff729ceaa 100644 --- a/crates/krishiv-session/src/observation.rs +++ b/crates/krishiv-session/src/observation.rs @@ -13,6 +13,13 @@ //! forever, as Spark's does; [`Observation::get_timeout`] bounds the wait. //! Before `observe` bound the handle it fails //! (`KRISHIV_NO_OBSERVE_BEFORE_GET`; PySpark's `NO_OBSERVE_BEFORE_GET`). +//! - **Blocking.** `get` waits on a condition variable: do not call it from async code on a +//! runtime thread (a current-thread runtime that also drives the action would deadlock); call +//! it from a plain thread, or after the action's future has been awaited. +//! - A failed finalization of the metrics (a metric that raises) is what `get` returns as its +//! error, as Spark's `get` rethrows it. +//! - **Counts follow Spark's** except where `docs/concepts/known-differences.md` lists a gap +//! (a `LIMIT` above a sort, an empty relation above the node). //! - The value is one row, a [`RecordBatch`] of one row whose columns are //! the metrics, named as written (`count(1)`, `max(v)` when not aliased); //! where the observed plan never ran its node (a `LIMIT 0` above it) the @@ -34,8 +41,8 @@ enum State { Unbound, /// Bound; no action has finished. Waiting, - /// The first action finished. - Done(RecordBatch), + /// The first action finished, with the metrics or the error finalizing them raised. + Done(Result), } /// What the frames and the handle share. @@ -72,13 +79,13 @@ impl Shared { /// An action finished: `row` is the metrics of the observed node, if it /// ran. Only the first action completes the observation. - pub(crate) fn complete(&self, row: Option<&RecordBatch>) { + pub(crate) fn complete(&self, row: Result, KrishivError>) { let mut state = self.lock(); if matches!(*state, State::Waiting) { - *state = State::Done( + *state = State::Done(row.map(|row| { row.cloned() - .unwrap_or_else(|| RecordBatch::new_empty(Arc::new(Schema::empty()))), - ); + .unwrap_or_else(|| RecordBatch::new_empty(Arc::new(Schema::empty()))) + })); self.done.notify_all(); } } @@ -120,10 +127,7 @@ impl Observation { /// An observation named by a fresh UUID (Spark's `Observation()`). #[must_use] pub fn new() -> Self { - // Unique within the process, which is all a plan's names need. - static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); - let n = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - Self::of(format!("observation-{n}")) + Self::of(crate::session::job_uuid()) } /// An observation named `name` (Spark's `Observation(name)`): the name @@ -164,7 +168,8 @@ impl Observation { /// without columns. /// /// # Errors - /// `KRISHIV_NO_OBSERVE_BEFORE_GET` when no frame was observed with it. + /// `KRISHIV_NO_OBSERVE_BEFORE_GET` when no frame was observed with it; the error that + /// finalizing the metrics raised. pub fn get(&self) -> Result { self.wait(None)? .ok_or_else(|| unreachable_wait(&self.shared.name)) @@ -190,7 +195,7 @@ impl Observation { } .into()); } - State::Done(row) => return Ok(Some(row.clone())), + State::Done(row) => return row.clone().map(Some), State::Waiting => {} } state = match deadline { diff --git a/crates/krishiv-session/src/result_cache.rs b/crates/krishiv-session/src/result_cache.rs index 3a2b75191..4ceb2e695 100644 --- a/crates/krishiv-session/src/result_cache.rs +++ b/crates/krishiv-session/src/result_cache.rs @@ -357,6 +357,9 @@ struct Found { /// implementer chooses, not the code). nondeterministic: bool, session_udf: bool, + /// The plan observes metrics (`Dataset.observe`): every action must run it, so its + /// observation completes. + observed: bool, } /// The functions that return a clock, the session or the environment: constant within a query, @@ -396,6 +399,9 @@ fn survey(plan: &Arc) -> Found { while let Some(step) = stack.pop() { match step { Step::Plan(p) => { + if matches!(&p.node, PlanNode::CollectMetrics { .. }) { + found.observed = true; + } if let PlanNode::Scan(s) = &p.node { match &s.source { ReadSource::Table { name, .. } => found.tables.push(name.clone()), @@ -438,6 +444,9 @@ impl Session { /// anything but permanent catalog tables, or calls a non-deterministic function. pub(crate) fn result_verdict(&self, a: &Analysis) -> Result { let found = survey(a.analyzed.plan.root()); + if found.observed { + return Err("it observes metrics"); + } if found.files { return Err("it reads files by path"); } diff --git a/crates/krishiv-session/src/session.rs b/crates/krishiv-session/src/session.rs index b41dffbdb..ee3832f78 100644 --- a/crates/krishiv-session/src/session.rs +++ b/crates/krishiv-session/src/session.rs @@ -1301,6 +1301,7 @@ impl Session { row: Arc::default(), identifier_spans: Arc::default(), identifier_written: Arc::default(), + observers: Arc::default(), cache_generation, cache_cap, sql_functions, @@ -1789,11 +1790,14 @@ impl Session { inner: stream, fold: pending.0.take(), }; - Ok(( + // The one place every run of a query passes (reads, writes, metrics, profiles): + // the observations of the frame complete when the stream ends, fails or drops. + let stream = crate::dataframe::deliver( Box::pin(stream.map(|b| b.and_then(into_record_batch))), - query, - row_id, - )) + Arc::clone(&query), + &env.observers, + ); + Ok((stream, query, row_id)) } } } diff --git a/crates/krishiv-session/src/session/env.rs b/crates/krishiv-session/src/session/env.rs index 2c78ab3ba..cfb9c442c 100644 --- a/crates/krishiv-session/src/session/env.rs +++ b/crates/krishiv-session/src/session/env.rs @@ -123,6 +123,9 @@ pub(crate) struct QueryEnv { pub(crate) identifier_spans: Arc<[krishiv_sql::ast::Span]>, /// The `IDENTIFIER(..)` text each of `identifier_spans` replaced (empty for `CACHE TABLE`). pub(crate) identifier_written: Arc<[String]>, + /// The observations and sinks of the frame the query runs (`Dataset.observe`): every run of + /// the query, whatever the action, completes them (`run_on_query_logged`). + pub(crate) observers: Arc, } /// A batch query's result rows: insert-only, so a weight other than 1 is a diff --git a/docs/reference/python/krishiv.sql.observation.md b/docs/reference/python/krishiv.sql.observation.md index 71e646de0..1ae57faf4 100644 --- a/docs/reference/python/krishiv.sql.observation.md +++ b/docs/reference/python/krishiv.sql.observation.md @@ -45,4 +45,5 @@ that never ran its metrics (`limit(0)` above it) gives `{}`. def name(self) -> str ``` -The name keying the metrics in the plan (not in PySpark's public API; kept for the Connect translation). +The name keying the metrics in the plan. A Krishiv extension: PySpark 4.2's `Observation` has only +`get`. diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 9634eee38..624018f65 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -185,6 +185,11 @@ message CollectMetricsNode { uint32 input = 1; string name = 2; uint32 measure = 3; + // The planner's rank of this node among the plan's observing nodes, in plan + // order: of several nodes with one name the highest is the observation's. + uint32 instance = 4; + // At most this many rows are folded (a `LIMIT` above takes no more). + optional uint64 limit = 5; } // feature: python_udf diff --git a/python/krishiv/sql/dataframe.py b/python/krishiv/sql/dataframe.py index 1b77a3e61..b047c18b1 100644 --- a/python/krishiv/sql/dataframe.py +++ b/python/krishiv/sql/dataframe.py @@ -28,7 +28,7 @@ from krishiv.sql.functions import first as _first from krishiv.sql.group import GroupedData from krishiv.sql.na import _NO_VALUE, DataFrameNaFunctions -from krishiv.sql.observation import Observation, _non_empty, _pyspark_error +from krishiv.sql.observation import Observation, _pyspark_error from krishiv.sql.profile_view import ProfileView from krishiv.sql.row import Row from krishiv.sql.stat import DataFrameStatFunctions @@ -440,10 +440,13 @@ def observe(self, observation: Observation | str, *exprs: Column) -> DataFrame: raise PySparkTypeError( errorClass="NOT_LIST_OF_COLUMN", messageParameters={"arg_name": "exprs", "arg_type": "tuple"} ) - if isinstance(observation, str): - _non_empty(observation) - observation = Observation(observation) try: + if isinstance(observation, str): + # The engine's analysis refuses an empty name (INVALID_OBSERVED_METRICS.MISSING_NAME, as + # classic Spark's string form does), so no client-side check here. + named = Observation.__new__(Observation) + named._native = _native.Observation(observation) + observation = named return self._new(self._df.observe(observation._native, [c._c for c in exprs])) except PySparkException as e: raise _pyspark_error(e) from e diff --git a/python/krishiv/sql/observation.py b/python/krishiv/sql/observation.py index 50b80a630..1805f6a9e 100644 --- a/python/krishiv/sql/observation.py +++ b/python/krishiv/sql/observation.py @@ -68,7 +68,8 @@ def __init__(self, name: str | None = None) -> None: @property def name(self) -> str: - """The name keying the metrics in the plan (not in PySpark's public API; kept for the Connect translation).""" + """The name keying the metrics in the plan. A Krishiv extension: PySpark 4.2's ``Observation`` has only + ``get``.""" return self._native.name @property diff --git a/python/public-api.txt b/python/public-api.txt index 45a27dfa0..f4109a0a0 100644 --- a/python/public-api.txt +++ b/python/public-api.txt @@ -63,6 +63,7 @@ DataFrame.def is_cached(self) -> bool DataFrame.def join(self, other: DataFrame, on: Column | None=..., using: list[str] | None=..., how: str=...) -> DataFrame DataFrame.def limit(self, n: int) -> DataFrame DataFrame.def local_checkpoint(self, eager: bool=..., level: tuple[bool, bool, bool, bool, int] | None=...) -> DataFrame +DataFrame.def observe(self, observation: Observation, metrics: list[Column]) -> DataFrame DataFrame.def offset(self, n: int) -> DataFrame DataFrame.def order_by(self, cols: list[Column]) -> DataFrame DataFrame.def persist(self, use_disk: bool, use_memory: bool, use_offheap: bool, deserialized: bool, replication: int) -> DataFrame @@ -82,6 +83,9 @@ DataFrame.def write(self, format: str, mode: str, path: str, options: list[tuple DataFrame.def write_table(self, insert_into: bool, table: str, format: str | None, mode: str, options: list[tuple[str, str]] | None=..., partition_by: list[str] | None=..., bucket_by: tuple[int, list[str]] | None=..., sort_by: list[str] | None=...) -> None DataFrame.def write_to_table(self, table: str, mode: str, provider: str | None=..., options: list[tuple[str, str]] | None=..., properties: list[tuple[str, str]] | None=..., partitioned_by: list[str] | None=..., cluster_by: list[str] | None=..., condition: str | None=...) -> None GroupedData.def agg(self, cols: list[Column]) -> DataFrame +Observation.@property def name(self) -> str +Observation.def __init__(self, name: str | None=...) -> None +Observation.def get(self) -> ArrowStream Session.@staticmethod def is_dataframe_table(name: str) -> bool Session.def __new__(cls, config: list[tuple[str, str]] | None=...) -> Session Session.def cache_reserved(self) -> int @@ -107,12 +111,13 @@ WindowSpec.def partition_by(self, cols: list[Column]) -> WindowSpec WindowSpec.def range_between(self, start: int, end: int) -> WindowSpec WindowSpec.def rows_between(self, start: int, end: int) -> WindowSpec __all__ -__all__ = ['AnonymousPythonUdf', 'ArrowStream', 'Column', 'DataFrame', 'GroupedData', 'Session', 'WindowSpec', '__allocator__', '__version__', '_udf_pool_shutdown', '_udf_runner_conf', 'call_function', 'col', 'count_distinct', 'count_star', 'expr', 'lit', 'map_literal'] +__all__ = ['AnonymousPythonUdf', 'ArrowStream', 'Column', 'DataFrame', 'GroupedData', 'Observation', 'Session', 'WindowSpec', '__allocator__', '__version__', '_udf_pool_shutdown', '_udf_runner_conf', 'call_function', 'col', 'count_distinct', 'count_star', 'expr', 'lit', 'map_literal'] class AnonymousPythonUdf() class ArrowStream() class Column() class DataFrame() class GroupedData() +class Observation() class Session() class WindowSpec() def call_function(name: str, args: list[Column]) -> Column @@ -193,7 +198,7 @@ def not_implemented(feature: str) -> PySparkNotImplementedError def unsupported_feature(feature: str) -> UnsupportedOperationException ## sql/__init__.py __all__ -__all__ = ['Catalog', 'Column', 'DataFrame', 'DataFrameNaFunctions', 'DataFrameReader', 'DataFrameStatFunctions', 'DataFrameWriter', 'DataFrameWriterV2', 'GroupedData', 'Row', 'RuntimeConfig', 'SparkSession', 'UDFRegistration', 'Window', 'WindowSpec', 'functions', 'types'] +__all__ = ['Catalog', 'Column', 'DataFrame', 'DataFrameNaFunctions', 'DataFrameReader', 'DataFrameStatFunctions', 'DataFrameWriter', 'DataFrameWriterV2', 'GroupedData', 'Observation', 'Row', 'RuntimeConfig', 'SparkSession', 'UDFRegistration', 'Window', 'WindowSpec', 'functions', 'types'] ## sql/_generated_functions.py __all__ __all__ = ['abs', 'acos', 'acosh', 'add_months', 'aes_decrypt', 'aes_encrypt', 'approx_count_distinct', 'approx_percentile', 'array', 'array_append', 'array_compact', 'array_contains', 'array_distinct', 'array_except', 'array_insert', 'array_intersect', 'array_join', 'array_max', 'array_min', 'array_position', 'array_prepend', 'array_remove', 'array_repeat', 'array_size', 'array_union', 'arrays_overlap', 'arrays_zip', 'ascii', 'asin', 'asinh', 'assert_true', 'atan', 'atan2', 'atanh', 'avg', 'base64', 'bin', 'bit_and', 'bit_count', 'bit_get', 'bit_length', 'bit_or', 'bit_xor', 'bitwiseNOT', 'bitwise_not', 'bool_and', 'bool_or', 'bround', 'btrim', 'cardinality', 'cbrt', 'ceil', 'ceiling', 'char', 'char_length', 'character_length', 'coalesce', 'collation', 'collect_list', 'collect_set', 'concat_ws', 'contains', 'conv', 'corr', 'cos', 'cosh', 'cot', 'count_if', 'covar_pop', 'covar_samp', 'crc32', 'create_map', 'csc', 'cume_dist', 'curdate', 'current_catalog', 'current_database', 'current_date', 'current_path', 'current_schema', 'current_timestamp', 'current_timezone', 'current_user', 'date_add', 'date_diff', 'date_format', 'date_from_unix_date', 'date_part', 'date_sub', 'date_trunc', 'dateadd', 'datediff', 'datepart', 'day', 'dayname', 'dayofmonth', 'dayofweek', 'dayofyear', 'decode', 'degrees', 'dense_rank', 'e', 'element_at', 'elt', 'encode', 'endswith', 'equal_null', 'exp', 'expm1', 'extract', 'factorial', 'find_in_set', 'first', 'first_value', 'flatten', 'floor', 'format_number', 'format_string', 'from_json', 'from_unixtime', 'from_utc_timestamp', 'get', 'get_json_object', 'getbit', 'greatest', 'hash', 'hex', 'hour', 'hypot', 'ifnull', 'initcap', 'input_file_block_length', 'input_file_block_start', 'input_file_name', 'instr', 'is_valid_utf8', 'is_variant_null', 'isnan', 'isnotnull', 'isnull', 'json_array_length', 'json_object_keys', 'json_tuple', 'kurtosis', 'lag', 'last', 'last_day', 'last_value', 'lcase', 'lead', 'least', 'left', 'length', 'levenshtein', 'listagg', 'ln', 'localtimestamp', 'locate', 'log10', 'log1p', 'log2', 'lower', 'lpad', 'ltrim', 'make_date', 'make_timestamp', 'make_timestamp_ltz', 'make_timestamp_ntz', 'make_valid_utf8', 'map_concat', 'map_contains_key', 'map_entries', 'map_from_arrays', 'map_from_entries', 'map_keys', 'map_values', 'mask', 'max', 'md5', 'mean', 'median', 'min', 'minute', 'mode', 'monotonically_increasing_id', 'month', 'monthname', 'months_between', 'named_struct', 'nanvl', 'negative', 'next_day', 'now', 'nth_value', 'ntile', 'nullif', 'nullifzero', 'nvl', 'nvl2', 'octet_length', 'overlay', 'parse_json', 'parse_url', 'percent_rank', 'percentile', 'percentile_approx', 'pi', 'pmod', 'position', 'positive', 'pow', 'printf', 'quarter', 'quote', 'radians', 'raise_error', 'rand', 'randn', 'randstr', 'rank', 'regexp', 'regexp_count', 'regexp_extract', 'regexp_extract_all', 'regexp_instr', 'regexp_like', 'regexp_replace', 'regexp_substr', 'regr_avgx', 'regr_avgy', 'regr_count', 'regr_intercept', 'regr_r2', 'regr_slope', 'regr_sxx', 'regr_sxy', 'regr_syy', 'repeat', 'replace', 'reverse', 'right', 'rint', 'round', 'row_number', 'rpad', 'rtrim', 'schema_of_json', 'schema_of_variant', 'sec', 'second', 'sequence', 'session_user', 'sha', 'sha1', 'sha2', 'shiftLeft', 'shiftRight', 'shiftRightUnsigned', 'shiftleft', 'shiftright', 'shiftrightunsigned', 'shuffle', 'sign', 'signum', 'sin', 'sinh', 'size', 'skewness', 'slice', 'sort_array', 'soundex', 'spark_partition_id', 'split', 'split_part', 'sqrt', 'startswith', 'std', 'stddev', 'stddev_pop', 'stddev_samp', 'str_to_map', 'struct', 'substr', 'substring', 'substring_index', 'sum', 'tan', 'tanh', 'time_bucket', 'timestamp_micros', 'timestamp_millis', 'timestamp_seconds', 'toDegrees', 'toRadians', 'to_binary', 'to_char', 'to_date', 'to_json', 'to_number', 'to_timestamp', 'to_timestamp_ltz', 'to_timestamp_ntz', 'to_unix_timestamp', 'to_utc_timestamp', 'to_varchar', 'to_variant_object', 'translate', 'trim', 'trunc', 'try_add', 'try_aes_decrypt', 'try_avg', 'try_divide', 'try_element_at', 'try_make_timestamp', 'try_make_timestamp_ltz', 'try_make_timestamp_ntz', 'try_mod', 'try_multiply', 'try_parse_json', 'try_parse_url', 'try_subtract', 'try_sum', 'try_to_binary', 'try_to_date', 'try_to_number', 'try_to_timestamp', 'try_url_decode', 'try_validate_utf8', 'try_variant_get', 'typeof', 'ucase', 'unbase64', 'unhex', 'uniform', 'unix_date', 'unix_micros', 'unix_millis', 'unix_seconds', 'upper', 'url_decode', 'url_encode', 'user', 'uuid', 'validate_utf8', 'var_pop', 'var_samp', 'variance', 'variant_get', 'version', 'weekday', 'weekofyear', 'width_bucket', 'xxhash64', 'year', 'zeroifnull'] @@ -1033,6 +1038,7 @@ DataFrame.def join(self, other: DataFrame, on: str | list[str] | Column | list[C DataFrame.def limit(self, num: int) -> DataFrame DataFrame.def localCheckpoint(self, eager: bool=True, storageLevel: StorageLevel | None=None) -> DataFrame DataFrame.def melt(self, ids: ColumnOrName | list[ColumnOrName] | tuple[ColumnOrName, ...], values: ColumnOrName | list[ColumnOrName] | tuple[ColumnOrName, ...] | None, variableColumnName: str, valueColumnName: str) -> DataFrame +DataFrame.def observe(self, observation: Observation | str, *exprs: Column) -> DataFrame DataFrame.def offset(self, num: int) -> DataFrame DataFrame.def orderBy(self, *cols: ColumnsOrNames, ascending: bool | list[bool]=True) -> DataFrame DataFrame.def persist(self, storageLevel: StorageLevel=StorageLevel.MEMORY_AND_DISK) -> DataFrame @@ -1111,6 +1117,14 @@ DataFrameNaFunctions.def drop(self, how: str | None='any', thresh: int | None=No DataFrameNaFunctions.def fill(self, value: Any, subset: Any=None) -> DataFrame DataFrameNaFunctions.def replace(self, to_replace: Any, value: Any=_NO_VALUE, subset: Any=None) -> DataFrame class DataFrameNaFunctions() +## sql/observation.py +Observation.@property def get(self) -> dict[str, Any] +Observation.@property def name(self) -> str +Observation.def __init__(self, name: str | None=None) -> None +Observation.def __repr__(self) -> str +__all__ +__all__ = ['Observation'] +class Observation() ## sql/pandas/functions.py ArrowUDFType.GROUPED_AGG ArrowUDFType.GROUPED_AGG_ITER From 114d2dd5d148dd2912a67296e142ac1c28a32914 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 05:48:49 +0000 Subject: [PATCH 07/15] tests: observe golden runner (classic and Connect legs) with KNOWN_GAPS, action and memory tests (P3-T084 review fixes) observe_goldens.rs runs every case of observe-goldens.json or lists it as client-only or a known gap that must still fail (sort+limit, ANSI sum overflow, scalar subquery). The Connect leg runs the flows_* and Connect-different cases through the server. Gaps are in docs/concepts/known-differences.md (follow-up P3-T084c). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-connect-server/tests/execute.rs | 287 +++++++ crates/krishiv-session/tests/observe.rs | 303 +++++-- .../krishiv-session/tests/observe_goldens.rs | 784 ++++++++++++++++++ docs/concepts/known-differences.md | 26 + python/tests/test_observe.py | 23 +- 5 files changed, 1333 insertions(+), 90 deletions(-) create mode 100644 crates/krishiv-session/tests/observe_goldens.rs diff --git a/crates/krishiv-connect-server/tests/execute.rs b/crates/krishiv-connect-server/tests/execute.rs index 11d73578e..b9a865400 100644 --- a/crates/krishiv-connect-server/tests/execute.rs +++ b/crates/krishiv-connect-server/tests/execute.rs @@ -1098,3 +1098,290 @@ async fn a_failed_query_sends_no_observed_metrics() { .expect_err("the table is missing"); assert_eq!(status.code(), Code::Internal); } + +// ---- the Connect leg of tests/fixtures/spark-4.2.0/observe-goldens.json (P3-T084) ---- +// +// Connect returns the metrics of the last action that ran the plan (no first-wins), all names +// in one `ExecutePlanResponse`, and nothing for a query that failed. The cases below are the +// fixture's Connect-only `flows_*` and the ones where the Connect leg differs from the classic +// one; the classic leg (`krishiv-session/tests/observe_goldens.rs`) covers the rest. Cases the +// engine does not match yet are asserted as what they give now (a gap), so a fix fails them. + +use pb::expression::literal::LiteralType; + +fn expression(e: pb::expression::ExprType) -> pb::Expression { + pb::Expression { + common: None, + expr_type: Some(e), + } +} + +fn attr(name: &str) -> pb::Expression { + expression(pb::expression::ExprType::UnresolvedAttribute( + pb::expression::UnresolvedAttribute { + unparsed_identifier: name.to_owned(), + ..Default::default() + }, + )) +} + +fn call(function: &str, arguments: Vec) -> pb::Expression { + expression(pb::expression::ExprType::UnresolvedFunction( + pb::expression::UnresolvedFunction { + function_name: function.to_owned(), + arguments, + ..Default::default() + }, + )) +} + +fn aliased(e: pb::Expression, name: &str) -> pb::Expression { + expression(pb::expression::ExprType::Alias(Box::new( + pb::expression::Alias { + expr: Some(Box::new(e)), + name: vec![name.to_owned()], + metadata: None, + }, + ))) +} + +fn base() -> pb::Relation { + sql( + "SELECT * FROM VALUES ('a', 1, 1.5D), ('b', 2, CAST(NULL AS DOUBLE)), \ + ('c', CAST(NULL AS INT), 3.0D) AS t(k, v, x)", + ) +} + +fn observe_as( + input: pb::Relation, + name: &str, + metrics: Vec, + plan_id: i64, +) -> pb::Relation { + pb::Relation { + common: Some(pb::RelationCommon { + plan_id: Some(plan_id), + ..Default::default() + }), + rel_type: Some(RelType::CollectMetrics(Box::new(pb::CollectMetrics { + input: Some(Box::new(input)), + name: name.to_owned(), + metrics, + }))), + } +} + +fn count_n() -> pb::Expression { + aliased(call("count", vec![literal_long(1)]), "n") +} + +fn limit_of(input: pb::Relation, limit: i32) -> pb::Relation { + relation(RelType::Limit(Box::new(pb::Limit { + input: Some(Box::new(input)), + limit, + }))) +} + +fn union_of(left: pb::Relation, right: pb::Relation) -> pb::Relation { + relation(RelType::SetOp(Box::new(pb::SetOperation { + left_input: Some(Box::new(left)), + right_input: Some(Box::new(right)), + set_op_type: pb::set_operation::SetOpType::Union as i32, + is_all: Some(true), + by_name: None, + allow_missing_columns: None, + }))) +} + +/// The observed metrics of a run: name -> (keys, values), and how many responses carried any. +/// Metric name -> (key, value) cells, and how many responses carried metrics. +type Observed = (BTreeMap)>>, usize); + +fn observed_of(responses: &[pb::ExecutePlanResponse]) -> Observed { + let carrying = responses + .iter() + .filter(|r| !r.observed_metrics.is_empty()) + .count(); + let all = responses + .iter() + .flat_map(|r| &r.observed_metrics) + .map(|m| { + let cells = m + .keys + .iter() + .cloned() + .zip(m.values.iter().map(|v| v.literal_type.clone())) + .collect(); + (m.name.clone(), cells) + }) + .collect(); + (all, carrying) +} + +fn long(v: i64) -> Option { + Some(LiteralType::Long(v)) +} + +async fn observed_run(plan: pb::Relation) -> Observed { + let run = start().await; + let responses = collect(&run.channel, request(root(plan))) + .await + .expect("execute"); + observed_of(&responses) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn connect_flows_string_observe_and_no_alias() { + // flows_string_observe: an unaliased literal is named by its text. + let plan = observe_as( + base(), + "nm", + vec![ + count_n(), + aliased(call("sum", vec![attr("v")]), "s"), + literal_long(7), + ], + 1, + ); + let (got, _) = observed_run(plan).await; + assert_eq!( + got["nm"], + vec![ + ("n".to_owned(), long(3)), + ("s".to_owned(), long(3)), + ("7".to_owned(), long(7)) + ] + ); + // flows_no_alias: metrics named as Spark names them. + let plan = observe_as( + base(), + "x", + vec![ + call("count", vec![literal_long(1)]), + call("max", vec![attr("v")]), + ], + 2, + ); + let (got, _) = observed_run(plan).await; + assert_eq!( + got["x"], + vec![ + ("count(1)".to_owned(), long(3)), + ("max(v)".to_owned(), Some(LiteralType::Integer(2))) + ] + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn connect_flows_two_nested_in_one_response() { + let inner = observe_as(base(), "a", vec![count_n()], 1); + let outer = observe_as( + inner, + "b", + vec![aliased(call("sum", vec![attr("v")]), "s")], + 2, + ); + let (got, carrying) = observed_run(outer).await; + assert_eq!(got["a"], vec![("n".to_owned(), long(3))]); + assert_eq!(got["b"], vec![("s".to_owned(), long(3))]); + assert_eq!(carrying, 1, "Spark sends all the names in one response"); +} + +/// Connect has no first-wins: each run reports its own metrics (`first_then_second_action`: +/// the count after a `limit(1)` run is 3), and the limit's run counts the rows it took. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn connect_each_action_reports_its_own_metrics() { + let run = start().await; + let observed = observe_as(base(), "m", vec![count_n()], 1); + let first = collect(&run.channel, request(root(limit_of(observed.clone(), 1)))) + .await + .expect("execute"); + assert_eq!(observed_of(&first).0["m"], vec![("n".to_owned(), long(1))]); + let second = collect(&run.channel, request(root(observed))) + .await + .expect("execute"); + assert_eq!(observed_of(&second).0["m"], vec![("n".to_owned(), long(3))]); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn connect_self_union_reports_one_instance_not_the_sum() { + let d = observe_as(base(), "m", vec![count_n()], 1); + let (got, _) = observed_run(union_of(d.clone(), d)).await; + assert_eq!(got["m"], vec![("n".to_owned(), long(3))]); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn connect_limit_zero_sends_no_metrics() { + let observed = observe_as(base(), "m", vec![count_n()], 1); + let (got, carrying) = observed_run(limit_of(observed, 0)).await; + assert!(got.is_empty(), "{got:?}"); + assert_eq!(carrying, 0); +} + +/// KNOWN GAP (docs/concepts/known-differences.md): Spark prunes an observe node under an empty +/// relation (`flows_observe_pruned_by_empty_relation` gives `e: {}`); the engine runs it and +/// reports the rows it read. When this fails the gap is closed: assert Spark's `{}`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn connect_gap_where_false_above_observe_is_not_pruned() { + let observed = observe_as(base(), "e", vec![count_n()], 1); + let filtered = relation(RelType::Filter(Box::new(pb::Filter { + input: Some(Box::new(observed)), + condition: Some(call("==", vec![literal_long(1), literal_long(0)])), + }))); + let (got, _) = observed_run(filtered).await; + assert_eq!(got["e"], vec![("n".to_owned(), long(3))]); +} + +/// KNOWN GAP: Spark supports a struct-valued metric (`types`); the server refuses it when it +/// builds the literal. The refusal names the feature, and the other values of the case work. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn connect_types_and_the_struct_metric_gap() { + let run = start().await; + let types = sql( + "SELECT CAST(1.10 AS DECIMAL(5,2)) d, CAST('NaN' AS DOUBLE) n, CAST(-0.0 AS DOUBLE) z, 1 i", + ); + let plan = observe_as( + types.clone(), + "m", + vec![ + aliased(call("sum", vec![attr("d")]), "sd"), + aliased(call("max", vec![attr("n")]), "mn"), + aliased(call("collect_list", vec![attr("i")]), "l"), + ], + 1, + ); + let responses = collect(&run.channel, request(root(plan))) + .await + .expect("execute"); + let (got, _) = observed_of(&responses); + let cells = &got["m"]; + assert_eq!(cells[0].0, "sd"); + assert!( + matches!(cells[0].1, Some(LiteralType::Decimal(_))), + "{cells:?}" + ); + assert!( + matches!(&cells[1].1, Some(LiteralType::Double(d)) if d.is_nan()), + "{cells:?}" + ); + assert!( + matches!(cells[2].1, Some(LiteralType::Array(_))), + "{cells:?}" + ); + let with_struct = observe_as( + types, + "m", + vec![aliased( + call( + "struct", + vec![call("max", vec![attr("i")]), call("min", vec![attr("i")])], + ), + "st", + )], + 1, + ); + let status = collect(&run.channel, request(root(with_struct))) + .await + .expect_err("a struct metric is not supported yet"); + assert_eq!(error_class(&status), "KRISHIV_UNSUPPORTED_FEATURE"); +} diff --git a/crates/krishiv-session/tests/observe.rs b/crates/krishiv-session/tests/observe.rs index 297ceea37..eb3276c30 100644 --- a/crates/krishiv-session/tests/observe.rs +++ b/crates/krishiv-session/tests/observe.rs @@ -1,21 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 -//! `DataFrame::observe` / `Observation` end to end (P3-T084): the classic -//! cases of `tests/fixtures/spark-4.2.0/observe-goldens.json` that do not -//! depend on row-at-a-time pulling, parallelism exactness, first action -//! wins, `LIMIT 0`, reuse errors and cancellation. +//! `DataFrame::observe` / `Observation` end to end (P3-T084): the actions that must complete an +//! observation (writes, table writes, metrics and profile runs, result-cache hits), a limit above +//! capping the count, the engine's own errors, memory and cancellation. The fixture's cases are +//! in `observe_goldens.rs`. #![cfg(test)] use std::sync::Arc; use std::time::Duration; use arrow_array::cast::AsArray; -use arrow_array::types::{Float64Type, Int32Type, Int64Type}; +use arrow_array::types::{Int32Type, Int64Type}; use arrow_array::{Array, Float64Array, Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; use futures::StreamExt; use futures::executor::block_on; -use krishiv_session::functions::{col, count, lit, max, sum}; -use krishiv_session::{Column, DataFrame, Observation, Session}; +use krishiv_session::functions::{col, count, expr, lit, max, sum}; +use krishiv_session::{DataFrame, Observation, Session}; use serde_json::Value; /// The generator's `base()`: `k string, v int, x double`. @@ -38,7 +38,7 @@ fn base(s: &Session) -> DataFrame { s.table("t") } -/// One metric cell as the fixture writes it: an integer or null. +/// One metric cell: an integer or null. fn cell(row: &RecordBatch, i: usize) -> Value { let c = row.column(i); if c.is_null(0) { @@ -47,11 +47,6 @@ fn cell(row: &RecordBatch, i: usize) -> Value { match c.data_type() { DataType::Int64 => Value::from(c.as_primitive::().value(0)), DataType::Int32 => Value::from(c.as_primitive::().value(0)), - DataType::Float64 => Value::from(c.as_primitive::().value(0)), - DataType::Utf8 | DataType::Utf8View => Value::from(c.as_string_opt::().map_or_else( - || c.as_string_view().value(0).to_owned(), - |s| s.value(0).to_owned(), - )), other => panic!("unexpected metric type {other}"), } } @@ -65,85 +60,12 @@ fn metrics(row: &RecordBatch) -> serde_json::Map { .collect() } -fn golden(case: &str) -> serde_json::Map { - let path = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../../tests/fixtures/spark-4.2.0/observe-goldens.json" - ); - let all: Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap(); - all["classic"][case]["metrics"].as_object().unwrap().clone() -} - /// Runs `df.count()` and returns what the observation holds. fn run(df: &DataFrame, obs: &Observation) -> serde_json::Map { block_on(df.count()).unwrap(); metrics(&obs.get().unwrap()) } -/// A golden case: its name, its metrics and what the frame is filtered by below the observation. -type Case = ( - &'static str, - Vec, - Box DataFrame>, -); - -#[test] -fn classic_goldens() { - let s = Session::new().unwrap(); - let t = base(&s); - let n = || count(lit(1)).alias("n"); - let cases: Vec = vec![ - ( - "basic_alias", - vec![ - count(lit(1)).alias("rows"), - sum(col("v")).alias("sv"), - max(col("k")).alias("mk"), - ], - Box::new(|d| d), - ), - ( - "literal_only", - vec![lit(42).alias("c"), lit("s").alias("s")], - Box::new(|d| d), - ), - ( - "filter_before_observe", - vec![n()], - Box::new(|d| d.filter(col("v").eq(1))), - ), - ( - "empty_input", - vec![ - count(lit(1)).alias("c"), - sum(col("v")).alias("s"), - max(col("v")).alias("mx"), - ], - Box::new(|d| d.filter(lit(1).eq(0))), - ), - ( - "nulls_all", - vec![ - sum(col("v")).alias("s"), - count(col("v")).alias("c"), - max(col("v")).alias("m"), - ], - Box::new(|d| d.filter(col("v").is_null())), - ), - ]; - for (name, metrics_of, below) in cases { - let obs = Observation::new(); - let df = below(t.clone()).observe(&obs, metrics_of).unwrap(); - let got = run(&df, &obs); - let want = golden(name); - // `empty_input` and `nulls_all` fixtures list more columns than the - // cases above build; every column built here must agree. - for (k, v) in &got { - assert_eq!(Some(v), want.get(k), "{name}: {k}"); - } - } -} - #[test] fn filter_above_counts_the_rows_before_it() { let s = Session::new().unwrap(); @@ -240,6 +162,9 @@ fn reuse_and_get_before_observe_fail() { ); } +/// What a dropped stream guarantees: the observation completes, and what it holds is +/// `{}` (the aborted drivers' partials had not arrived) or a count no larger than the input. +/// Spark also reports a partial value for a cancelled action. #[test] fn a_cancelled_action_completes_the_observation() { let s = Session::new().unwrap(); @@ -257,9 +182,209 @@ fn a_cancelled_action_completes_the_observation() { .get_timeout(Duration::from_secs(30)) .unwrap() .expect("completes after the stream is dropped"); - // Whatever flowed before the cancellation was counted, no more than the input. - let n = metrics(&got).get("n").and_then(Value::as_i64).unwrap_or(0); - assert!((0..=50_000_000).contains(&n)); + match metrics(&got).get("n") { + None => assert_eq!(got.num_columns(), 0), + Some(n) => assert!((0..=50_000_000).contains(&n.as_i64().unwrap())), + } +} + +/// Every way an action runs the plan completes the observation (one hook, +/// `Session::run_on_query_logged`), not only `execute_stream`. +#[test] +fn a_parquet_write_completes_the_observation() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s) + .observe(&obs, [count(lit(1)).alias("n"), sum(col("v")).alias("s")]) + .unwrap(); + let dir = tempfile::tempdir().unwrap(); + block_on(df.write().parquet(dir.path().join("o").to_str().unwrap())).unwrap(); + let got = metrics( + &obs.get_timeout(Duration::from_secs(10)) + .unwrap() + .expect("completes"), + ); + assert_eq!((got["n"].as_i64(), got["s"].as_i64()), (Some(3), Some(3))); +} + +#[test] +fn a_table_write_completes_the_observation() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(df.write().save_as_table("observed_out")).unwrap(); + let got = obs + .get_timeout(Duration::from_secs(10)) + .unwrap() + .expect("completes"); + assert_eq!(metrics(&got)["n"], 3); + // Appending is another write through the same hook. + let again = Observation::new(); + let df = base(&s) + .observe(&again, [count(lit(1)).alias("n")]) + .unwrap(); + block_on(df.write().mode("append").insert_into("observed_out")).unwrap(); + let got = again + .get_timeout(Duration::from_secs(10)) + .unwrap() + .expect("completes"); + assert_eq!(metrics(&got)["n"], 3); +} + +#[test] +fn collect_with_metrics_and_profile_complete_the_observation() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(df.collect_with_metrics()).unwrap(); + let got = obs + .get_timeout(Duration::from_secs(10)) + .unwrap() + .expect("completes"); + assert_eq!(metrics(&got)["n"], 3); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(df.explain(krishiv_session::ExplainMode::Simple)).unwrap(); + assert!( + obs.get_timeout(Duration::from_millis(200)) + .unwrap() + .is_none() + ); + block_on(df.collect_with_profile()).unwrap(); + let got = obs + .get_timeout(Duration::from_secs(10)) + .unwrap() + .expect("completes"); + assert_eq!(metrics(&got)["n"], 3); +} + +/// A plan that observes is never served from the result cache: every action runs it, so every +/// observation completes. +#[test] +fn a_result_cache_hit_does_not_skip_the_observation() { + let s = Session::new().unwrap(); + s.conf().set("krishiv.resultCache.enabled", "true").unwrap(); + block_on( + block_on(s.sql( + "CREATE TABLE rc_observed USING parquet AS SELECT 1 AS a UNION ALL SELECT 2 AS a", + )) + .unwrap() + .collect(), + ) + .unwrap(); + for _ in 0..3 { + let obs = Observation::named("m").unwrap(); + let df = block_on(s.sql("SELECT a FROM rc_observed")) + .unwrap() + .observe(&obs, [count(lit(1)).alias("n")]) + .unwrap(); + assert_eq!(block_on(df.count()).unwrap(), 2); + let got = obs + .get_timeout(Duration::from_secs(10)) + .unwrap() + .expect("completes"); + assert_eq!(metrics(&got)["n"], 2); + } +} + +/// Spark counts the rows its consumer pulled: a limit above (through row-for-row operators) +/// takes no more than that, not the batch it read. +#[test] +fn a_limit_above_caps_the_count_at_the_rows_it_takes() { + let s = Session::new().unwrap(); + let count_of = |df: &DataFrame, limit: i32| { + let obs = Observation::new(); + let observed = df.observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(observed.limit(limit).collect()).unwrap(); + metrics(&obs.get_timeout(Duration::from_secs(10)).unwrap().unwrap())["n"] + .as_i64() + .unwrap() + }; + assert_eq!(count_of(&base(&s), 1), 1); + assert_eq!(count_of(&base(&s), 2), 2); + assert_eq!(count_of(&base(&s), 10), 3); + assert_eq!(count_of(&s.range(0, 1000, 1, Some(4)), 4), 4); + // A projection between passes every row on. + let obs = Observation::new(); + let df = s + .range(0, 1000, 1, Some(4)) + .observe(&obs, [count(lit(1)).alias("n")]) + .unwrap() + .select([col("id")]); + block_on(df.limit(7).collect()).unwrap(); + assert_eq!(metrics(&obs.get().unwrap())["n"], 7); + // A filter between changes which rows the limit pulled: the count is whole batches at least. + let obs = Observation::new(); + let df = base(&s) + .observe(&obs, [count(lit(1)).alias("n")]) + .unwrap() + .filter(col("v").is_not_null()); + block_on(df.limit(1).collect()).unwrap(); + assert!(metrics(&obs.get().unwrap())["n"].as_i64().unwrap() >= 1); +} + +/// A scalar subquery in a metric is refused as an unsupported feature, never an internal error. +#[test] +fn a_scalar_subquery_metric_is_unsupported_not_internal() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s) + .observe( + &obs, + [expr("(select max(id) from range(3))").unwrap().alias("s")], + ) + .unwrap(); + let e = block_on(df.count()).unwrap_err(); + assert_eq!(e.error_class().name(), "KRISHIV_UNSUPPORTED_FEATURE", "{e}"); +} + +/// The state of a metric is charged to the query's pool: an unbounded `collect_list` fails with +/// `UNABLE_TO_ACQUIRE_MEMORY` under a tiny limit and runs under a sufficient one. +#[test] +fn metric_state_is_charged_to_the_memory_pool() { + let tiny = Session::builder() + .config("krishiv.memory.limit", "64kb") + .build() + .unwrap(); + let obs = Observation::new(); + let df = tiny + .range(0, 1_000_000, 1, Some(2)) + .observe(&obs, [expr("collect_list(id)").unwrap().alias("l")]) + .unwrap(); + let e = block_on(df.count()).unwrap_err(); + assert_eq!(e.error_class().name(), "UNABLE_TO_ACQUIRE_MEMORY", "{e}"); + let roomy = Session::builder() + .config("krishiv.memory.limit", "512mb") + .build() + .unwrap(); + let obs = Observation::new(); + let df = roomy + .range(0, 100_000, 1, Some(2)) + .observe(&obs, [expr("size(collect_list(id))").unwrap().alias("l")]) + .unwrap(); + assert_eq!(block_on(df.count()).unwrap(), 100_000); + assert_eq!(metrics(&obs.get().unwrap())["l"], 100_000); +} + +/// A metric that raises while the states are finalized is what `get` returns (Spark rethrows the +/// root cause), not an empty row. +#[test] +fn a_failed_finalization_is_what_get_returns() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + // The state folds; the metric over the results divides by zero when finalized. + let df = base(&s) + .observe( + &obs, + [expr("count(1) / (count(1) - 3)").unwrap().alias("r")], + ) + .unwrap(); + s.conf().set("spark.sql.ansi.enabled", "true").unwrap(); + block_on(df.count()).unwrap(); + let e = obs + .get_timeout(Duration::from_secs(10)) + .expect_err("the finalization error, not a silent {}"); + assert_eq!(e.error_class().name(), "DIVIDE_BY_ZERO", "{e}"); } #[test] diff --git a/crates/krishiv-session/tests/observe_goldens.rs b/crates/krishiv-session/tests/observe_goldens.rs new file mode 100644 index 000000000..d80d11642 --- /dev/null +++ b/crates/krishiv-session/tests/observe_goldens.rs @@ -0,0 +1,784 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `Dataset.observe` / `Observation` against the classic leg of +//! `tests/fixtures/spark-4.2.0/observe-goldens.json` (real Spark 4.2.0 runs, generated by +//! `tests/scripts/gen_observe_goldens.py`; P3-T084). Every fixture case is either run here +//! (`cases`), run by the Python tests because the client raises it (`CLIENT_ONLY`), or listed in +//! `KNOWN_GAPS` with the reason it is not matched yet: a gap's case must still fail, so a fix +//! removes it from the list (the `window_range_month_goldens.rs` pattern). The Connect leg runs +//! through the server in `krishiv-connect-server/tests/execute.rs`. +#![cfg(test)] + +use std::sync::Arc; +use std::time::Duration; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Decimal128Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type}; +use arrow_array::{Array, ArrayRef, Float64Array, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::executor::block_on; +use krishiv_common::error::KrishivError; +use krishiv_session::functions::{col, count, expr, lit, max, sum}; +use krishiv_session::{Column, DataFrame, ExplainMode, JoinType, Observation, Session}; +use serde_json::{Map, Value, json}; + +/// Cases the Python client raises before the engine is reached (`python/tests/test_observe.py`) +/// or that describe a PySpark class's attributes. +const CLIENT_ONLY: &[(&str, &str)] = &[ + ( + "anonymous_observation", + "reads the Python Observation's private name", + ), + ("python_attrs", "attributes of PySpark's Observation class"), + ( + "err_empty_observation_name", + "VALUE_NOT_NON_EMPTY_STR is the Python client's", + ), + ( + "err_get_before_observe", + "NO_OBSERVE_BEFORE_GET is the Python client's (the engine's KRISHIV_NO_OBSERVE_BEFORE_GET is tested in observe.rs)", + ), + ( + "err_reuse", + "REUSE_OBSERVATION is the Python client's (KRISHIV_OBSERVATION_REUSED is tested in observe.rs)", + ), + ("err_no_exprs", "CANNOT_BE_EMPTY is the Python client's"), + ( + "err_no_exprs_string", + "CANNOT_BE_EMPTY is the Python client's", + ), + ( + "err_observation_name_type", + "NOT_EXPECTED_TYPE is the Python client's", + ), + ( + "limit_zero", + "classic Spark fails in the JVM here (an AssertionError in toPyRow); the Connect leg gives {}, run through the server", + ), +]; + +/// Cases the engine does not match yet. Each must still fail; remove it when it passes. +const KNOWN_GAPS: &[(&str, &str)] = &[ + ( + "observation_through_sort_limit_topn", + "Spark's TakeOrderedAndProject pulls 20 of the 1000 rows (a partition-dependent count); the engine's top-N reads all 1000 (P3-T084c)", + ), + ( + "err_scalar_subquery", + "Spark runs a scalar subquery in a metric; the engine refuses it with KRISHIV_UNSUPPORTED_FEATURE (P3-T084c)", + ), + ( + "ansi_on_overflow_sum", + "Spark's metric aggregate saturates a long sum at 9223372036854775807 under ANSI; the engine raises ARITHMETIC_OVERFLOW (P3-T084c)", + ), +]; + +/// Cases whose `metrics` are compared as a subset of the fixture's columns (the case leaves out +/// what depends on which of the parallel drivers' rows a `LIMIT` takes). +const SUBSET_METRICS: &[&str] = &["limit_big_range"]; + +fn fixture() -> Value { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../tests/fixtures/spark-4.2.0/observe-goldens.json" + ); + serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap() +} + +/// A session over the generator's `base()`: `t(k string, v int, x double)`, ANSI on. +fn session() -> Session { + let s = Session::new().unwrap(); + s.conf().set("spark.sql.ansi.enabled", "true").unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Utf8, true), + Field::new("v", DataType::Int32, true), + Field::new("x", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(1), Some(2), None])), + Arc::new(Float64Array::from(vec![Some(1.5), None, Some(3.0)])), + ], + ) + .unwrap(); + s.register_batches("t", schema, vec![batch]).unwrap(); + s +} + +fn base(s: &Session) -> DataFrame { + s.table("t") +} + +fn sql(s: &Session, text: &str) -> DataFrame { + block_on(s.sql(text)).unwrap() +} + +fn n() -> Column { + count(lit(1)).alias("n") +} + +fn e(text: &str) -> Column { + expr(text).unwrap() +} + +/// Python's `repr` of a float. +fn py_float(f: f64) -> String { + if f.is_nan() { + "nan".to_owned() + } else if f.is_infinite() { + if f > 0.0 { "inf" } else { "-inf" }.to_owned() + } else if f == f.trunc() && f.abs() < 1e16 { + format!("{f:.1}") + } else { + format!("{f}") + } +} + +/// One cell as the fixture writes values (`norm` in the generator). +fn cell(c: &ArrayRef, i: usize) -> Value { + if c.is_null(i) { + return Value::Null; + } + match c.data_type() { + DataType::Boolean => Value::from(c.as_boolean().value(i)), + DataType::Int8 => Value::from(c.as_primitive::().value(i)), + DataType::Int16 => Value::from(c.as_primitive::().value(i)), + DataType::Int32 => Value::from(c.as_primitive::().value(i)), + DataType::Int64 => Value::from(c.as_primitive::().value(i)), + DataType::Float64 => Value::from(py_float(c.as_primitive::().value(i))), + DataType::Utf8 => Value::from(c.as_string::().value(i)), + DataType::Utf8View => Value::from(c.as_string_view().value(i)), + DataType::Decimal128(_, scale) => { + let raw = c.as_primitive::().value(i); + let scale = usize::try_from(*scale).unwrap(); + let digits = raw.unsigned_abs().to_string(); + let digits = format!("{digits:0>width$}", width = scale + 1); + let (int, frac) = digits.split_at(digits.len() - scale); + let sign = if raw < 0 { "-" } else { "" }; + let repr = if scale == 0 { + format!("{sign}{int}") + } else { + format!("{sign}{int}.{frac}") + }; + json!({"py_type": "Decimal", "repr": repr}) + } + DataType::List(_) => { + let list = c.as_list::().value(i); + Value::Array((0..list.len()).map(|j| cell(&list, j)).collect()) + } + DataType::Struct(_) => { + let st = c.as_struct(); + Value::Array( + (0..st.num_columns()) + .map(|j| cell(st.column(j), i)) + .collect(), + ) + } + other => panic!("unexpected metric type {other}"), + } +} + +/// The metrics row as a dict: a repeated name keeps its last column (Python's dict). +fn metrics(row: &RecordBatch) -> Value { + let mut out = Map::new(); + for (i, f) in row.schema().fields().iter().enumerate() { + out.insert(f.name().clone(), cell(row.column(i), 0)); + } + Value::Object(out) +} + +/// What `Observation.get` returns (`getres` in the generator), or `blocked` when no action has +/// finished. +fn got(o: &Observation) -> Value { + match o.get_timeout(Duration::from_millis(400)) { + Ok(Some(row)) => json!({"get": "ok", "metrics": metrics(&row)}), + Ok(None) => json!({"get": "blocked"}), + Err(e) => json!({"get": "error", "error": error(&e)}), + } +} + +fn error(e: &KrishivError) -> Value { + json!({ + "errorClass": e.error_class().name(), + "sqlState": e.sql_state(), + "outcome": "error", + }) +} + +#[expect( + clippy::needless_pass_by_value, + reason = "builders pass frames and values on" +)] +fn merged(mut a: Value, b: Value) -> Value { + a.as_object_mut() + .unwrap() + .extend(b.as_object().unwrap().clone()); + a +} + +fn counted(df: &DataFrame) -> Value { + // An action's failure is the case's value, so the comparison (or a gap) reports it. + match block_on(df.count()) { + Ok(n) => json!(n), + Err(e) => error(&e), + } +} + +/// `df.observe(o, metrics)`, the action, and what `get` holds (`run` in the generator). +#[expect( + clippy::needless_pass_by_value, + reason = "builders pass frames and values on" +)] +fn run(df: DataFrame, metrics_of: Vec, action: impl FnOnce(&DataFrame) -> Value) -> Value { + let o = Observation::named("m").unwrap(); + let observed = df.observe(&o, metrics_of).unwrap(); + let result = action(&observed); + let row = if result.is_null() { + json!({}) + } else { + json!({"action_result": result}) + }; + merged(row, got(&o)) +} + +fn count_action(d: &DataFrame) -> Value { + counted(d) +} + +/// An analysis or run error of `exprs` over `df`, or what `get` holds if it ran. +#[expect( + clippy::needless_pass_by_value, + reason = "builders pass frames and values on" +)] +fn err_case(df: DataFrame, exprs: Vec) -> Value { + let o = Observation::named("m").unwrap(); + let observed = df.observe(&o, exprs).unwrap(); + match block_on(observed.count()) { + Err(e) => error(&e), + Ok(_) => merged(json!({"outcome": "ok_unexpected"}), got(&o)), + } +} + +type Case = fn(&Session) -> Value; + +#[expect(clippy::too_many_lines, reason = "one table of the fixture's cases")] +fn cases() -> Vec<(&'static str, Case)> { + vec![ + ("basic_alias", |s| { + run( + base(s), + vec![ + count(lit(1)).alias("rows"), + sum(col("v")).alias("sv"), + max(col("k")).alias("mk"), + ], + count_action, + ) + }), + ("no_alias_names", |s| { + run( + base(s), + vec![ + count(lit(1)), + max(col("v")), + sum(col("v") + lit(1)), + lit(42), + ], + count_action, + ) + }), + ("literal_only", |s| { + run( + base(s), + vec![lit(42).alias("c"), lit("s").alias("s")], + count_action, + ) + }), + ("composite_expr", |s| { + run( + base(s), + vec![ + (sum(col("v")) + e("avg(x)") - lit(1)).alias("e"), + e("count(*)").alias("n"), + ], + count_action, + ) + }), + ("empty_input", |s| { + run( + base(s).filter(lit(1).eq(0)), + vec![ + n().alias("c"), + sum(col("v")).alias("s"), + max(col("k")).alias("mx"), + e("avg(x)").alias("a"), + e("collect_list(k)").alias("l"), + ], + count_action, + ) + }), + ("empty_input_no_rows_collected", |s| { + run(base(s).filter(lit(1).eq(0)), vec![n().alias("c")], |d| { + let rows: usize = block_on(d.collect()) + .unwrap() + .iter() + .map(RecordBatch::num_rows) + .sum(); + assert_eq!(rows, 0); + json!([]) + }) + }), + ("types", |s| { + run( + sql( + s, + "select cast(1.10 as decimal(5,2)) d, cast('NaN' as double) n, cast(-0.0 as double) z, 1 i", + ), + vec![ + e("sum(d)").alias("sd"), + e("max(n)").alias("mn"), + e("min(z)").alias("mz"), + e("sum(i)").alias("si"), + e("avg(i)").alias("ai"), + e("collect_list(i)").alias("l"), + e("struct(max(i), min(i))").alias("st"), + ], + count_action, + ) + }), + ("nulls_all", |s| { + run( + sql(s, "select cast(null as int) v"), + vec![ + sum(col("v")).alias("s"), + count(col("v")).alias("c"), + max(col("v")).alias("m"), + ], + count_action, + ) + }), + ("aggregate_with_alias_inside_nested_alias", |s| { + run( + base(s), + vec![sum(col("v")).alias("a").alias("b")], + count_action, + ) + }), + ("observed_metric_column_name_collision", |s| { + run(base(s), vec![n(), sum(col("v")).alias("n")], count_action) + }), + ("count_action_prunes_columns", |s| { + run(base(s), vec![sum(col("v")).alias("s")], count_action) + }), + ("filter_after_observe", |s| { + run(base(s), vec![n()], |d| counted(&d.filter(col("v").eq(1)))) + }), + ("filter_before_observe", |s| { + let o = Observation::named("m").unwrap(); + let df = base(s).filter(col("v").eq(1)).observe(&o, [n()]).unwrap(); + block_on(df.collect()).unwrap(); + got(&o) + }), + ("exact_under_parallelism", |s| { + run( + s.range(0, 100_000, 1, Some(8)), + vec![n(), sum(col("id")).alias("s"), max(col("id")).alias("mx")], + count_action, + ) + }), + ("conditional_metric_no_raise_unselected", |s| { + run( + sql(s, "select * from values (0),(2) as t(v)"), + vec![e("sum(case when v != 0 then 10 / v end)").alias("s")], + count_action, + ) + }), + ("ansi_on_overflow_sum", |s| { + let text = "select cast(9223372036854775807 as bigint) v union all select 1"; + match block_on(s.sql(text)) { + Err(e) => json!({"stage": "sql", "error": error(&e)}), + Ok(df) => run(df, vec![sum(col("v")).alias("s")], count_action), + } + }), + ("ansi_off_overflow_sum", |s| { + s.conf().set("spark.sql.ansi.enabled", "false").unwrap(); + run( + sql( + s, + "select cast(9223372036854775807 as bigint) v union all select 1", + ), + vec![ + sum(col("v")).alias("s"), + (max(col("v")) + lit(1)).alias("p1"), + ], + count_action, + ) + }), + ("ansi_on_cast_error_inside_metric", |s| { + err_case(base(s), vec![e("sum(cast(k as int))").alias("s")]) + }), + // limits: the rows a limit above takes are the rows counted + ("limit_after_observe", |s| { + run(base(s), vec![n()], |d| { + block_on(d.limit(1).collect()).unwrap(); + Value::Null + }) + }), + ("take_after_observe", |s| { + run(base(s), vec![n()], |d| { + block_on(d.limit(1).collect()).unwrap(); + Value::Null + }) + }), + ("limit_big_range", |s| { + let mut v = run( + s.range(0, 100_000, 1, Some(4)), + vec![n(), sum(col("id")).alias("s")], + |d| { + block_on(d.limit(3).collect()).unwrap(); + Value::Null + }, + ); + // `s` sums the rows the racing drivers handed the limit (Spark: partition 0's rows). + v["metrics"].as_object_mut().unwrap().remove("s"); + v + }), + ("show_action_limit20", |s| { + run(s.range(0, 1000, 1, Some(4)), vec![n()], |d| { + // `show(3)` reads `limit(4)`: one more row than it shows. + block_on(d.limit(4).collect()).unwrap(); + Value::Null + }) + }), + ("observation_through_sort_limit_topn", |s| { + run(s.range(0, 1000, 1, Some(4)), vec![n()], |d| { + block_on(d.order_by([col("id")]).limit(5).collect()).unwrap(); + Value::Null + }) + }), + ("first_then_second_action", |s| { + let o = Observation::named("m").unwrap(); + let df = base(s).observe(&o, [n()]).unwrap(); + block_on(df.limit(1).collect()).unwrap(); + let first = got(&o)["metrics"].clone(); + block_on(df.count()).unwrap(); + let second = got(&o)["metrics"].clone(); + json!({"after_limit1": first, "after_count": second}) + }), + ("multiple_actions_first_wins", |s| { + let o = Observation::named("m").unwrap(); + let df = base(s) + .observe(&o, [n(), sum(col("v")).alias("s")]) + .unwrap(); + let c1 = block_on(df.count()).unwrap(); + let after_first = got(&o); + let c2 = block_on(df.filter(col("v").is_not_null()).count()).unwrap(); + let c3: usize = block_on(df.collect()) + .unwrap() + .iter() + .map(RecordBatch::num_rows) + .sum(); + json!({"counts": [c1, c2, c3], "after_first": after_first, "after_more": got(&o)}) + }), + ("get_blocks_without_action", |s| { + let o = Observation::named("m").unwrap(); + base(s).observe(&o, [n()]).unwrap(); + got(&o) + }), + ("get_after_transformation_only_no_action_then_action", |s| { + let o = Observation::named("m").unwrap(); + let df = base(s).observe(&o, [n()]).unwrap(); + let before = got(&o); + block_on(df.count()).unwrap(); + json!({"before_action": before, "after_action": got(&o)}) + }), + ("explain_does_not_trigger", |s| { + let o = Observation::named("m").unwrap(); + let df = base(s).observe(&o, [n()]).unwrap(); + block_on(df.explain(ExplainMode::Simple)).unwrap(); + got(&o) + }), + ("observation_on_action_collect_schema_only", |s| { + let o = Observation::named("m").unwrap(); + let df = base(s).observe(&o, [n()]).unwrap(); + block_on(df.schema()).unwrap(); + json!({"get_after_schema": got(&o)}) + }), + // caches + ("cache", |s| { + let o = Observation::named("m").unwrap(); + let df = base(s).observe(&o, [n()]).unwrap().cache(); + let a = block_on(df.count()).unwrap(); + let after_first = got(&o); + let b = block_on(df.count()).unwrap(); + json!({"counts": [a, b], "after_first": after_first, "after_second": got(&o)}) + }), + ("cache_then_new_observe_on_cached", |s| { + let df = base(s).cache(); + block_on(df.count()).unwrap(); + let o = Observation::named("m").unwrap(); + block_on(df.observe(&o, [n()]).unwrap().collect()).unwrap(); + got(&o) + }), + // several nodes in one plan + ("union_two_branches_distinct_names", |s| { + let (oa, ob) = ( + Observation::named("a").unwrap(), + Observation::named("b").unwrap(), + ); + let a = base(s).observe(&oa, [n()]).unwrap(); + let b = base(s).filter(col("v").eq(1)).observe(&ob, [n()]).unwrap(); + json!({"count": counted(&a.union(&b)), "a": got(&oa), "b": got(&ob)}) + }), + ("union_two_branches_same_name_diff_df", |s| { + let a = base(s).observe_named("dup", [n()]); + let b = base(s).observe_named("dup", [n()]); + match block_on(a.union(&b).count()) { + Err(e) => error(&e), + Ok(_) => json!({"outcome": "ok_unexpected"}), + } + }), + ("self_join_same_df", |s| { + let o = Observation::named("m").unwrap(); + let d = base(s).observe(&o, [n()]).unwrap(); + let l = d.alias("l"); + let r = d.alias("r"); + let c = counted(&l.join(&r, l.col("k").eq(r.col("k")), JoinType::Inner)); + merged(json!({"count": c}), got(&o)) + }), + ("self_union_same_df", |s| { + let o = Observation::named("m").unwrap(); + let d = base(s).observe(&o, [n()]).unwrap(); + merged(json!({"count": counted(&d.union(&d))}), got(&o)) + }), + ("join_two_observed", |s| { + let (oa, ob) = ( + Observation::named("a").unwrap(), + Observation::named("b").unwrap(), + ); + let a = base(s).observe(&oa, [n()]).unwrap(); + let b = base(s).observe(&ob, [sum(col("v")).alias("s")]).unwrap(); + let c = counted(&a.join_using(&b, ["k"], JoinType::Inner)); + json!({"count": c, "a": got(&oa), "b": got(&ob)}) + }), + ("same_name_string_same_df_in_both_branches", |s| { + let d = base(s).observe_named("dup", [n()]); + json!({"count": counted(&d.union(&d))}) + }), + ("nested_observe_different_names", |s| { + let (oa, ob) = ( + Observation::named("a").unwrap(), + Observation::named("b").unwrap(), + ); + let d = base(s) + .observe(&oa, [n()]) + .unwrap() + .filter(col("v").eq(1)) + .observe(&ob, [count(lit(1)).alias("m")]) + .unwrap(); + json!({"count": counted(&d), "a": got(&oa), "b": got(&ob)}) + }), + ("nested_observe_same_name_diff_df", |s| { + let d = base(s) + .observe_named("dup", [n()]) + .filter(col("v").eq(1)) + .observe_named("dup", [n()]); + match block_on(d.count()) { + Err(e) => error(&e), + Ok(_) => json!({"outcome": "ok_unexpected"}), + } + }), + ("observation_in_groupby_downstream", |s| { + let o = Observation::named("m").unwrap(); + let d = base(s) + .observe(&o, [n()]) + .unwrap() + .group_by([col("k")]) + .agg([count(lit(1))]); + block_on(d.collect()).unwrap(); + got(&o) + }), + ( + "string_observe_returns_df_only", + |s| json!({"count": counted(&base(s).observe_named("nm", [n()]))}), + ), + // failures + ("failed_query", |s| { + let o = Observation::named("m").unwrap(); + let d = s + .range(0, 5, 1, None) + .observe(&o, [n()]) + .unwrap() + .select_expr(["1/(id-2) as r"]) + .unwrap(); + let action = match block_on(d.collect()) { + Err(e) => error(&e), + Ok(_) => json!("ok_unexpected"), + }; + json!({"action": {"errorClass": action["errorClass"], "sqlState": action["sqlState"]}, "get": got(&o)["get"]}) + }), + ("failed_query_then_reuse_df_succeeds", |s| { + let o = Observation::named("m").unwrap(); + let d = s.range(0, 5, 1, None).observe(&o, [n()]).unwrap(); + drop(block_on( + d.select_expr(["1/(id-2) as r"]).unwrap().collect(), + )); + let after_fail = got(&o)["get"].clone(); + block_on(d.count()).unwrap(); + json!({"after_fail": {"get": after_fail}, "after_success": {"get": got(&o)["get"]}}) + }), + ("write_action", |s| { + let o = Observation::named("m").unwrap(); + let d = base(s) + .observe(&o, [n(), sum(col("v")).alias("s")]) + .unwrap(); + let dir = tempfile::tempdir().unwrap(); + block_on(d.write().parquet(dir.path().join("o").to_str().unwrap())).unwrap(); + got(&o) + }), + // metric analysis + ("err_attribute", |s| err_case(base(s), vec![col("v")])), + ("err_attribute_alias_mixed", |s| { + err_case(base(s), vec![n().alias("c"), col("v").alias("raw")]) + }), + ("err_nondeterministic_rand", |s| { + err_case(base(s), vec![e("rand()").alias("r")]) + }), + ("ok_nondeterministic_inside_agg", |s| { + let v = err_case(base(s), vec![e("sum(rand())").alias("r")]); + json!({"outcome": v["outcome"], "get": v["get"]}) + }), + ("err_nondeterministic_plus_agg", |s| { + err_case(base(s), vec![(n() + e("rand()")).alias("r")]) + }), + ("err_nested_aggregate", |s| { + err_case(base(s), vec![e("sum(max(v))").alias("n")]) + }), + ("err_distinct", |s| { + err_case(base(s), vec![e("count(distinct v)").alias("d")]) + }), + ("err_distinct_sum", |s| { + err_case(base(s), vec![e("sum(distinct v)").alias("d")]) + }), + ("err_filter", |s| { + err_case(base(s), vec![e("count(v) filter (where v > 1)").alias("f")]) + }), + ("err_window", |s| { + err_case(base(s), vec![e("sum(v) over (order by k)").alias("w")]) + }), + ("err_generator", |s| { + err_case(base(s), vec![e("explode(array(1, 2))").alias("g")]) + }), + ("err_unresolved_column", |s| { + err_case(base(s), vec![e("sum(nope)").alias("s")]) + }), + ("err_scalar_subquery", |s| { + err_case( + base(s), + vec![e("(select max(id) from range(3))").alias("s")], + ) + }), + ("ok_aggregate_over_subquery_free_udf_like", |s| { + err_case(base(s), vec![e("sum(length(k))").alias("s")]) + }), + ("err_unresolved_star", |s| { + err_case(base(s), vec![e("count(*)").alias("c"), col("*")]) + }), + ("ok_approx_count_distinct", |s| { + err_case(base(s), vec![e("approx_count_distinct(v)").alias("a")]) + }), + ("err_first_ignorenulls_ok", |s| { + err_case(base(s), vec![e("first(v, true)").alias("f")]) + }), + ("ok_count_if", |s| { + err_case(base(s), vec![e("count_if(v > 0)").alias("c")]) + }), + ("err_empty_name_string_api", |s| { + let d = base(s).observe_named("", [n()]); + match block_on(d.count()) { + Err(e) => error(&e), + Ok(_) => json!({"outcome": "ok_unexpected"}), + } + }), + ] +} + +/// `got` against `want`: every key `got` has must match (so a case compares the columns it +/// builds); a metrics dict must match key for key unless the case is in `SUBSET_METRICS`. +fn matches(got: &Value, want: &Value, subset: bool, key: &str) -> bool { + match (got, want) { + (Value::Object(g), Value::Object(w)) => { + let exact = key == "metrics" && !subset; + (!exact || g.len() == w.len()) + && g.iter() + .all(|(k, v)| w.get(k).is_some_and(|x| matches(v, x, subset, k))) + } + (Value::Array(g), Value::Array(w)) => { + g.len() == w.len() && g.iter().zip(w).all(|(a, b)| matches(a, b, subset, key)) + } + _ => got == want, + } +} + +fn outcome(name: &str, case: Case, want: &Value) -> Result<(), String> { + // A panic inside a case (an unexpected `unwrap`) is that case's failure, with its name. + let got = + std::panic::catch_unwind(|| case(&session())).map_err(|_| format!("{name}: panicked"))?; + let subset = SUBSET_METRICS.contains(&name); + if matches(&got, want, subset, "") { + Ok(()) + } else { + Err(format!("{name}: got {got}\n want {want}")) + } +} + +#[test] +fn every_fixture_case_is_run_skipped_or_a_known_gap() { + let all = fixture(); + let classic = all["classic"].as_object().unwrap(); + let cases = cases(); + let mut seen: Vec<&str> = cases.iter().map(|(n, _)| *n).collect(); + seen.extend(CLIENT_ONLY.iter().map(|(n, _)| *n)); + let missing: Vec<&String> = classic + .keys() + .filter(|k| !seen.contains(&k.as_str())) + .collect(); + assert!( + missing.is_empty(), + "fixture cases nothing runs: {missing:?}" + ); + for name in seen { + assert!(classic.contains_key(name), "{name} is not in the fixture"); + } + for (name, _) in KNOWN_GAPS { + assert!( + cases.iter().any(|(n, _)| n == name), + "gap {name} is not a case" + ); + } +} + +#[test] +fn classic_goldens_match_spark() { + let all = fixture(); + let classic = all["classic"].as_object().unwrap(); + let mut failures = Vec::new(); + let mut fixed = Vec::new(); + for (name, case) in cases() { + let result = outcome(name, case, &classic[name]); + let gap = KNOWN_GAPS.iter().any(|(g, _)| *g == name); + match (result, gap) { + (Err(why), false) => failures.push(why), + (Ok(()), true) => fixed.push(name), + // What the gap looks like now (visible with `--nocapture`). + (Err(why), true) => eprintln!("known gap: {why}"), + _ => {} + } + } + assert!(failures.is_empty(), "{}", failures.join("\n")); + assert!( + fixed.is_empty(), + "these known gaps now match Spark: remove them from KNOWN_GAPS: {fixed:?}" + ); +} diff --git a/docs/concepts/known-differences.md b/docs/concepts/known-differences.md index 6ae90ef7e..dc35065d9 100644 --- a/docs/concepts/known-differences.md +++ b/docs/concepts/known-differences.md @@ -225,3 +225,29 @@ do not see it, as in Spark. `system.caches` lists the session's live checkpoints stream. Krishiv's replay buffer and resume come in a later release. - **ReleaseExecute without a release type** fails with `CONNECT_INVALID_PLAN.INVALID_ONE_OF_FIELD_NOT_SET`; Spark throws an `UnsupportedOperationException` without an error class. + +## `Dataset.observe`: where the counts differ from Spark's (follow-up P3-T084c) + +`DataFrame::observe` / PySpark `Observation` match Spark 4.2.0 on every case of +`tests/fixtures/spark-4.2.0/observe-goldens.json` except the ones below +(`crates/krishiv-session/tests/observe_goldens.rs` pins them in `KNOWN_GAPS`, each must still fail there, so a +fix removes the line). A `LIMIT` (and so `take`, `first`, `show`) directly above the observed node, +or above projections only, counts exactly the rows it takes, as Spark's `CollectMetricsExec` counts the rows +its consumer pulled; the cap is one budget per node instance, so over several partitions the rows +counted are the first the racing drivers hand over (Spark's are partition 0's: a `sum` metric can differ). + +- **Sort and limit:** `df.observe(..).orderBy(..).limit(n)` reads every row in the engine (`n` = the input); + Spark's `TakeOrderedAndProject` pulls a partition-dependent subset (20 of 1000 in the golden). +- **ANSI `sum` overflow:** Spark's metric aggregate returns the saturated `9223372036854775807` for a long + `sum` that overflows under ANSI; the engine raises `ARITHMETIC_OVERFLOW`. With ANSI off both wrap. +- **Scalar subquery in a metric:** Spark runs it; the engine refuses it with `KRISHIV_UNSUPPORTED_FEATURE`. +- **Empty relation above:** `observe(..).where(false)` (Spark prunes the node; `{}`) counts the rows read. + Only `LIMIT 0` prunes. Streaming frames are not refused as Spark refuses them. +- **Connect:** struct, map and interval metric values are refused when the server builds the + `ObservedMetrics` literal (after the rows were sent); a name with no plan id reports `plan_id` 0. +- **Cancelled actions:** a dropped result stream completes the observation with `{}` or the rows + folded before the drivers stopped (Spark: a partial value); a failed finalization of a metric is what `get` + raises. +- **`Observation.get` blocks** a thread; do not call it from async code on a runtime thread. +- **Result cache:** a plan that observes is never served from the result cache (every action runs it). + diff --git a/python/tests/test_observe.py b/python/tests/test_observe.py index e440737f9..968f3f59e 100644 --- a/python/tests/test_observe.py +++ b/python/tests/test_observe.py @@ -4,10 +4,11 @@ from __future__ import annotations import threading +from pathlib import Path import pytest -from krishiv.errors import PySparkAssertionError, PySparkTypeError, PySparkValueError +from krishiv.errors import PySparkAssertionError, PySparkException, PySparkTypeError, PySparkValueError from krishiv.sql import DataFrame, Observation, SparkSession from krishiv.sql import functions as F @@ -86,6 +87,26 @@ def test_argument_errors(df: DataFrame) -> None: assert c2.value.getCondition() == "CANNOT_BE_EMPTY" +def test_empty_string_name_is_refused_by_the_engine(df: DataFrame) -> None: + # classic Spark's string form: INVALID_OBSERVED_METRICS.MISSING_NAME (the class API's VALUE_NOT_NON_EMPTY_STR + # is the Observation constructor's) + with pytest.raises(PySparkException) as e: + df.observe("", F.count(F.lit(1)).alias("n")).count() + assert e.value.getCondition() == "INVALID_OBSERVED_METRICS.MISSING_NAME" + + +def test_limit_above_counts_the_rows_it_takes(df: DataFrame) -> None: + obs = Observation() + df.observe(obs, F.count(F.lit(1)).alias("n")).limit(1).collect() + assert obs.get == {"n": 1} + + +def test_write_completes_the_observation(df: DataFrame, tmp_path: Path) -> None: + obs = Observation() + df.observe(obs, F.count(F.lit(1)).alias("n")).write.parquet(str(tmp_path / "o")) + assert obs.get == {"n": 10} + + def test_limit_zero_gives_empty(df: DataFrame) -> None: obs = Observation() df.observe(obs, F.count(F.lit(1)).alias("n")).limit(0).collect() From 219d23b9bd271f94f9e0526be79f4a6e303e4891 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 10:56:25 +0000 Subject: [PATCH 08/15] python, docs: Observation stub __new__, regenerated migration guide, changelog fragment (#756) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- docs/guides/pyspark-migration.md | 1 + python/krishiv/_native.pyi | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/guides/pyspark-migration.md b/docs/guides/pyspark-migration.md index 9e227450e..ed833e6ea 100644 --- a/docs/guides/pyspark-migration.md +++ b/docs/guides/pyspark-migration.md @@ -163,6 +163,7 @@ Each row links to the full description in [known differences](../concepts/known- | [The relation cache is per session](../concepts/known-differences.md#the-relation-cache-is-per-session) | | [Checkpoint files and where a checkpoint lives](../concepts/known-differences.md#checkpoint-files-and-where-a-checkpoint-lives) | | [`krishiv server` (Spark Connect)](../concepts/known-differences.md#krishiv-server-spark-connect) | +| [`Dataset.observe`: where the counts differ from Spark's (follow-up P3-T084c)](../concepts/known-differences.md#datasetobserve-where-the-counts-differ-from-sparks-follow-up-p3-t084c) | ## Where it runs diff --git a/python/krishiv/_native.pyi b/python/krishiv/_native.pyi index 92e0a587f..75ebc55e4 100644 --- a/python/krishiv/_native.pyi +++ b/python/krishiv/_native.pyi @@ -69,7 +69,7 @@ class Session: @final class Observation: - def __init__(self, name: str | None = ...) -> None: ... + def __new__(cls, name: str | None = ...) -> Observation: ... @property def name(self) -> str: ... def get(self) -> ArrowStream: ... From 34272296f38559dc030b61111ff6c9f316007d3c Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 10:56:48 +0000 Subject: [PATCH 09/15] changelog: fragment for #756 Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- changelog.d/756.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/756.added.md diff --git a/changelog.d/756.added.md b/changelog.d/756.added.md new file mode 100644 index 000000000..326da2856 --- /dev/null +++ b/changelog.d/756.added.md @@ -0,0 +1 @@ +Observations (DataFrame.observe, Observation) are supported: metrics are computed in the query's own plan. From a1ff8ec57e09b330f881bc923468af06763960a5 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 11:42:20 +0000 Subject: [PATCH 10/15] plan, exec, session, connect: observe review fixes (prefix-cache names, memory charge, per-observation errors, temp views) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-connect-server/src/execute.rs | 17 ++++- crates/krishiv-exec/src/aggregate/observe.rs | 17 ++++- crates/krishiv-exec/src/context.rs | 46 +++++++++++-- crates/krishiv-plan/src/analyzer/prefix.rs | 18 +++++ crates/krishiv-plan/src/tests/prefix_cache.rs | 38 ++++++++++- crates/krishiv-session/src/catalog.rs | 53 +++++++++++++++ crates/krishiv-session/src/dataframe.rs | 13 +++- .../krishiv-session/src/dataframe/observe.rs | 61 ++++++++++------- crates/krishiv-session/tests/observe.rs | 66 +++++++++++++++++++ 9 files changed, 295 insertions(+), 34 deletions(-) diff --git a/crates/krishiv-connect-server/src/execute.rs b/crates/krishiv-connect-server/src/execute.rs index 20a996ee1..074d564da 100644 --- a/crates/krishiv-connect-server/src/execute.rs +++ b/crates/krishiv-connect-server/src/execute.rs @@ -672,16 +672,29 @@ async fn produce( sender.send(metrics).await?; // The observed nodes' metrics follow the SQL metrics, all in one response as Spark sends // them; a query that ended in an error never gets here, so it sends none. + // An observation whose metrics failed to finalize fails the action (Spark: the metric's error + // is the action's), after the others' metrics were sent. let mut all = Vec::new(); + let mut failed = None; for (name, row) in sink.take().unwrap_or_default() { - let plan_id = plan_ids.get(&name).copied().unwrap_or(0); - all.push(observed::observed_metrics(&name, plan_id, &row)?); + match row { + Ok(row) => { + let plan_id = plan_ids.get(&name).copied().unwrap_or(0); + all.push(observed::observed_metrics(&name, plan_id, &row)?); + } + Err(e) => { + failed.get_or_insert(e); + } + } } if !all.is_empty() { let mut response = sender.response(None); response.observed_metrics = all; sender.send(response).await?; } + if let Some(e) = failed { + return Err(e.into()); + } // The end is atomic with interrupts: once finished, an interrupt // interrupts nothing; an interrupt that came first turns the end into // OPERATION_CANCELED (Spark's `ExecuteThreadRunner`, v4.3.0 lines diff --git a/crates/krishiv-exec/src/aggregate/observe.rs b/crates/krishiv-exec/src/aggregate/observe.rs index 9debe8cef..8b056c376 100644 --- a/crates/krishiv-exec/src/aggregate/observe.rs +++ b/crates/krishiv-exec/src/aggregate/observe.rs @@ -68,11 +68,21 @@ impl Finalizer { /// The metrics' row over `partials`, the states of every driver of the /// instance: the aggregates' zero state if there are none. - pub(crate) fn finalize(&self, partials: &[RecordBatch]) -> Result { + /// + /// The merge table is charged to `reservation` as it grows. + pub(crate) fn finalize( + &self, + partials: &[RecordBatch], + reservation: &mut Reservation, + ) -> Result { let mut table = AggTable::new(&self.core, &HashSeed::new(None))?; for states in partials { let p = table.prepare(&Batch::new(states.clone()), true)?; + let want = table.projected(&p); + reservation.try_grow(want.saturating_sub(reservation.size()))?; table.apply(&p)?; + let over = table.memory_size().saturating_sub(reservation.size()); + reservation.try_grow(over)?; } let results = table.output_batch(&[0])?; let out = self.post.project(&Batch::new(results), 0)?; @@ -162,8 +172,9 @@ impl Local { return Ok(()); } let states = self.table.state_batch(&[0])?; - self.query.observe_partial(&self.finalizer, states); - Ok(()) + // The states move to the query's own charge: release the driver's first. + self.reservation.shrink(self.reservation.size()); + self.query.observe_partial(&self.finalizer, states) } } diff --git a/crates/krishiv-exec/src/context.rs b/crates/krishiv-exec/src/context.rs index 835f872b0..b37c9fff4 100644 --- a/crates/krishiv-exec/src/context.rs +++ b/crates/krishiv-exec/src/context.rs @@ -229,6 +229,9 @@ pub struct QueryContext { /// The partial states of each observing node (`Dataset.observe`), by name /// and the planner's instance number (see [`Self::observe_partial`]). observed: Mutex>, + /// The memory consumer that holds the retained partial states of `observed` (charged to the + /// pool, so a variable-size metric such as `collect_list` cannot outgrow the query's limit). + observed_memory: Mutex>, finished: AtomicBool, finished_event: Event, /// Blocking closures in flight; the query ends only once they have. @@ -277,6 +280,7 @@ impl QueryContext { spills: Mutex::new(BTreeMap::new()), scan_gate: OnceLock::new(), observed: Mutex::new(BTreeMap::new()), + observed_memory: Mutex::new(None), finished: AtomicBool::new(false), finished_event: Event::default(), blocking: AtomicUsize::new(0), @@ -286,12 +290,22 @@ impl QueryContext { /// A driver's partial states of an observation: its task metric, sent /// when the driver's input ended (or was abandoned). The states of one - /// instance of the node merge in [`Self::observed`]. + /// instance of the node merge in [`Self::observed`]. The retained states + /// are charged to the query's pool. + /// + /// # Errors + /// `UNABLE_TO_ACQUIRE_MEMORY` when the pool cannot hold them. pub(crate) fn observe_partial( &self, finalizer: &Arc, states: arrow_array::RecordBatch, - ) { + ) -> Result<()> { + let bytes = states.get_array_memory_size(); + { + let mut held = lock(&self.observed_memory); + let reservation = held.get_or_insert_with(|| self.observed_consumer().reservation()); + reservation.try_grow(bytes)?; + } lock(&self.observed) .entry((finalizer.name().to_owned(), finalizer.instance())) .or_insert_with(|| Observed { @@ -300,6 +314,11 @@ impl QueryContext { }) .partials .push(states); + Ok(()) + } + + fn observed_consumer(&self) -> Arc { + ConsumerHandle::register(Arc::clone(&self.pool), "observed-metrics", None) } /// The observations (`Dataset.observe`) of the query so far: each name's @@ -310,15 +329,34 @@ impl QueryContext { /// Spark. /// /// # Errors - /// A metric's own error (finalizing runs its expressions). + /// A metric's own error (finalizing runs its expressions), or + /// `UNABLE_TO_ACQUIRE_MEMORY` (the first one, by name). #[doc(hidden)] pub fn observed(&self) -> Result> { + self.observed_each() + .into_iter() + .map(|(n, r)| Ok((n, r?))) + .collect() + } + + /// [`Self::observed`] with each observation's own outcome: one whose + /// finalization fails does not take the others with it. + #[doc(hidden)] + pub fn observed_each(&self) -> BTreeMap> { let all = lock(&self.observed); // Keys sort by name then instance: the last one of a name overwrites. let last: BTreeMap<&String, &Observed> = all.iter().map(|((name, _), o)| (name, o)).collect(); + let consumer = self.observed_consumer(); last.into_iter() - .map(|(name, o)| Ok((name.clone(), o.finalizer.finalize(&o.partials)?))) + .map(|(name, o)| { + // The merge table is charged while it lives. + let mut reservation = consumer.reservation(); + ( + name.clone(), + o.finalizer.finalize(&o.partials, &mut reservation), + ) + }) .collect() } diff --git a/crates/krishiv-plan/src/analyzer/prefix.rs b/crates/krishiv-plan/src/analyzer/prefix.rs index 8484bbcbc..cb72796bf 100644 --- a/crates/krishiv-plan/src/analyzer/prefix.rs +++ b/crates/krishiv-plan/src/analyzer/prefix.rs @@ -46,6 +46,9 @@ pub(super) struct PrefixEntry { /// The tables the root reads: as written, the metadata resolved, and the securable index. tables: Vec<(Vec, Arc, usize)>, sql_calls: usize, + /// The observation names the root registered (`Dataset.observe`), with the identity key of the + /// relation that holds each: a warm run must still refuse a second relation using the name. + observed: Vec<(String, u128)>, } impl Resolver<'_> { @@ -99,6 +102,11 @@ impl Resolver<'_> { columns: self.columns.clone(), tables, sql_calls: self.sql_calls, + observed: self + .observed + .iter() + .filter_map(|(name, a)| Some((name.clone(), *self.keys.get(a)?))) + .collect(), }) } @@ -157,6 +165,16 @@ impl Resolver<'_> { self.closed.insert((at, env.ctes_id), r.clone()); } } + for (name, key) in &entry.observed { + let holder = self + .keys + .iter() + .filter(|(_, k)| *k == key) + .find_map(|(a, _)| self.canon.get(a).map(|rep| addr(&**rep))); + if let Some(at) = holder { + self.observed.insert(name.clone(), at); + } + } self.warm_started = true; self.next_id = entry.end_id; self.declared = entry.declared.clone(); diff --git a/crates/krishiv-plan/src/tests/prefix_cache.rs b/crates/krishiv-plan/src/tests/prefix_cache.rs index ed45f52b8..1a03c5318 100644 --- a/crates/krishiv-plan/src/tests/prefix_cache.rs +++ b/crates/krishiv-plan/src/tests/prefix_cache.rs @@ -13,7 +13,7 @@ use crate::{ ColumnRef, Expr, ExprNode, JoinCondition, JoinKind, LogicalPlan, Origin, Plan, PlanId, PlanNode, }; -use super::dsl::{add, int}; +use super::dsl::{add, count_star, int}; fn table(cols: &[&str]) -> MemoryCatalog { let cols = cols @@ -267,3 +267,39 @@ fn a_moved_version_is_a_miss() { "version 8 is a miss" ); } + +/// A prefix-cache hit keeps the observation names the prefix registered: a second relation may not +/// reuse one (`DUPLICATED_METRICS_NAME`), warm as cold. +#[test] +fn a_warm_start_keeps_observed_names() { + let c = ctx(&["a", "b", "c"]); + let cache = CachingAnalyzer::new(); + let metric = || Arc::new(count_star().alias("n")); + let observe = |input: LogicalPlan, id: i64| { + LogicalPlan::new(PlanNode::CollectMetrics { + input: Arc::new(input), + name: "dup".to_owned(), + metrics: vec![metric()], + }) + .with_origin(Origin::plan(PlanId(id))) + }; + let first = observe( + LogicalPlan::table(["t"]).with_origin(Origin::plan(PlanId(0))), + 1, + ); + assert!(!same(&cache, first.clone(), &c)); + let other = observe( + LogicalPlan::table(["u"]).with_origin(Origin::plan(PlanId(5))), + 6, + ); + let both = LogicalPlan::new(PlanNode::Join { + left: Arc::new(first), + right: Arc::new(other), + kind: JoinKind::Cross, + condition: JoinCondition::None, + }) + .with_origin(Origin::plan(PlanId(9))); + let p = Plan::try_new(both).unwrap(); + let e = block_on(cache.analyze(&p, &c)).expect_err("two relations, one name"); + assert_eq!(e.error.error_class().name(), "DUPLICATED_METRICS_NAME"); +} diff --git a/crates/krishiv-session/src/catalog.rs b/crates/krishiv-session/src/catalog.rs index 7d5722515..db0df3305 100644 --- a/crates/krishiv-session/src/catalog.rs +++ b/crates/krishiv-session/src/catalog.rs @@ -145,6 +145,9 @@ pub(crate) fn is_reserved_view_name(name: &str) -> bool { #[derive(Debug, Default)] pub(crate) struct SessionCatalog { tables: RwLock)>>, + /// What a view's frame kept (its `Observation` handles), by the view's key: a query that reads + /// the view completes them. Cleared when the entry is replaced or dropped. + view_keeps: RwLock>>>, epoch: AtomicU64, /// Hidden names handed out for snapshot pins (see [`SessionCatalog::put_view`]). hidden: AtomicU64, @@ -168,6 +171,7 @@ impl SessionCatalog { pub(crate) fn over(store: Arc, warehouse: Arc) -> Self { Self { tables: RwLock::default(), + view_keeps: RwLock::default(), epoch: AtomicU64::new(0), hidden: AtomicU64::new(0), store, @@ -201,6 +205,12 @@ impl SessionCatalog { .clone(); Self { tables: RwLock::new(tables), + view_keeps: RwLock::new( + self.view_keeps + .read() + .unwrap_or_else(PoisonError::into_inner) + .clone(), + ), epoch: AtomicU64::new(0), hidden: AtomicU64::new(self.hidden.load(Ordering::Acquire)), store: Arc::clone(&self.store), @@ -214,6 +224,7 @@ impl SessionCatalog { /// owned that no other live view still owns. pub(crate) fn put(&self, name: &str, table: TempTable) { let mut tables = self.tables.write().unwrap_or_else(PoisonError::into_inner); + self.forget_keeps(&name.to_lowercase()); tables.insert(name.to_lowercase(), (name.to_owned(), Arc::new(table))); Self::collect_garbage(&mut tables); self.epoch.fetch_add(1, Ordering::AcqRel); @@ -233,6 +244,7 @@ impl SessionCatalog { new_pins: Vec<(String, Arc)>, ) { let mut tables = self.tables.write().unwrap_or_else(PoisonError::into_inner); + self.forget_keeps(&name.to_lowercase()); for (hidden, table) in new_pins { tables.insert(hidden.clone(), (hidden, table)); } @@ -252,6 +264,46 @@ impl SessionCatalog { self.epoch.fetch_add(1, Ordering::AcqRel); } + fn forget_keeps(&self, key: &str) { + self.view_keeps + .write() + .unwrap_or_else(PoisonError::into_inner) + .remove(key); + } + + /// Records what the frame that made view `name` kept (its `Observation` handles), for the + /// queries that read the view. + pub(crate) fn set_view_keeps( + &self, + name: &str, + keep: Vec>, + ) { + if keep.is_empty() { + return; + } + self.view_keeps + .write() + .unwrap_or_else(PoisonError::into_inner) + .insert(name.to_lowercase(), keep); + } + + /// What the view read as `parts` (a temporary view, or `global_temp.v`) kept. + pub(crate) fn view_keeps(&self, parts: &[String]) -> Vec> { + let key = match parts { + [view] => view.to_lowercase(), + [ns, view] if ns.eq_ignore_ascii_case(GLOBAL_NAMESPACE) => { + global_key(view).to_lowercase() + } + _ => return Vec::new(), + }; + self.view_keeps + .read() + .unwrap_or_else(PoisonError::into_inner) + .get(&key) + .cloned() + .unwrap_or_default() + } + /// Removes the hidden entries nothing needs; the caller holds the write lock and bumps the epoch. fn collect_garbage(tables: &mut BTreeMap)>) { let mut live = BTreeSet::new(); @@ -284,6 +336,7 @@ impl SessionCatalog { } let mut tables = self.tables.write().unwrap_or_else(PoisonError::into_inner); let removed = tables.remove(&name.to_lowercase()).is_some(); + self.forget_keeps(&name.to_lowercase()); if removed { Self::collect_garbage(&mut tables); self.epoch.fetch_add(1, Ordering::AcqRel); diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index cf7d4ee64..bd700f4c7 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -1053,7 +1053,13 @@ impl DataFrame { if let Some(t) = text { env.set_text(t); } - env.observers = Arc::new(self.observers()); + let through_views = analyzed + .views + .iter() + .filter(|v| v.security == krishiv_plan::analyzer::views::ViewSecurity::Temporary) + .flat_map(|v| self.session.state.catalog.view_keeps(&v.name)) + .collect(); + env.observers = Arc::new(self.observers_through(through_views)); let analysis = Arc::new(Analysis { env: Arc::new(env), analyzed: Arc::new(analyzed), @@ -1808,6 +1814,11 @@ impl DataFrame { .state .catalog .put_view(name, expanded, columns, pins, new_pins); + // A query that reads the view runs its observed nodes: it completes the same observations. + self.session + .state + .catalog + .set_view_keeps(name, self.keep.clone()); self.session.state.cache.forget_table(name); Ok(()) } diff --git a/crates/krishiv-session/src/dataframe/observe.rs b/crates/krishiv-session/src/dataframe/observe.rs index 79174bdac..11bb38ec9 100644 --- a/crates/krishiv-session/src/dataframe/observe.rs +++ b/crates/krishiv-session/src/dataframe/observe.rs @@ -76,8 +76,12 @@ impl DataFrame { out } - /// The observations and sinks this frame carries, for the query's env. - pub(super) fn observers(&self) -> Observers { + /// The observations and sinks this frame carries, and those of the temporary views it reads + /// (`through_views`), for the query's env. + pub(super) fn observers_through( + &self, + through_views: Vec>, + ) -> Observers { let mut out = Observers::default(); for k in &self.keep { if let Ok(shared) = Arc::clone(k).downcast::() { @@ -86,6 +90,14 @@ impl DataFrame { out.sinks.push(sink); } } + // A view's own sinks belong to the frame that made it, not to this query. + for k in through_views { + if let Ok(shared) = k.downcast::() + && !out.observations.iter().any(|o| Arc::ptr_eq(o, &shared)) + { + out.observations.push(shared); + } + } out } } @@ -98,12 +110,13 @@ pub(crate) struct Observers { } /// The state of an [`ObservedSink`]. -pub(crate) type SinkSlot = Mutex>>; +pub(crate) type SinkSlot = Mutex>>>; /// Where an action leaves the metrics of the observed nodes of its plan, by -/// name (one row each): set when the result was read to its end without an -/// error, replaced by each later action (Spark Connect: the last action -/// wins, a failed query sends none). For the Connect server. +/// name (one row each, or the error finalizing that observation raised): +/// set when the result was read to its end without an error, replaced by +/// each later action (Spark Connect: the last action wins, a failed query +/// sends none). For the Connect server. #[doc(hidden)] #[derive(Clone, Default)] pub struct ObservedSink { @@ -124,9 +137,10 @@ impl ObservedSink { } /// The metrics the last successful action left, once; `None` when no - /// action has finished without an error. + /// action has finished without an error. An observation whose own + /// metrics failed to finalize is an `Err` beside the others' rows. #[must_use] - pub fn take(&self) -> Option> { + pub fn take(&self) -> Option>> { self.slot .lock() .unwrap_or_else(PoisonError::into_inner) @@ -170,25 +184,26 @@ impl Delivering { return; }; // On the carriers a caller reads (decimals, intervals), as the action's rows are. - let observed = self.query.observed().and_then(|all| { - all.into_iter() - .map(|(name, row)| Ok((name, crate::carrier::client_batch(&row)?))) - .collect::, KrishivError>>() - }); - if self.finished - && let Ok(all) = &observed - { + let observed: BTreeMap> = self + .query + .observed_each() + .into_iter() + .map(|(name, row)| { + let row = row.and_then(|row| crate::carrier::client_batch(&row)); + (name, row) + }) + .collect(); + if self.finished { for sink in sinks { - *sink.lock().unwrap_or_else(PoisonError::into_inner) = Some(all.clone()); + *sink.lock().unwrap_or_else(PoisonError::into_inner) = Some(observed.clone()); } } for o in observations { - o.complete( - observed - .as_ref() - .map(|all| all.get(o.name())) - .map_err(Clone::clone), - ); + o.complete(match observed.get(o.name()) { + Some(Ok(row)) => Ok(Some(row)), + Some(Err(e)) => Err(e.clone()), + None => Ok(None), + }); } } } diff --git a/crates/krishiv-session/tests/observe.rs b/crates/krishiv-session/tests/observe.rs index eb3276c30..26f1fe5b0 100644 --- a/crates/krishiv-session/tests/observe.rs +++ b/crates/krishiv-session/tests/observe.rs @@ -398,3 +398,69 @@ fn invalid_metrics_fail_the_action_with_spark_classes() { "INVALID_OBSERVED_METRICS.NON_AGGREGATE_FUNC_ARG_IS_ATTRIBUTE" ); } + +/// The states a driver retained after it ended and the table that merges them are charged to the +/// pool too: the drivers' own states fit, their merge does not. +#[test] +fn retained_states_and_the_merge_are_charged_to_the_memory_pool() { + let s = Session::builder() + .config("krishiv.memory.limit", "2200kb") + .build() + .unwrap(); + let obs = Observation::new(); + let df = s + .range(0, 200_000, 1, Some(2)) + .observe(&obs, [expr("collect_list(id)").unwrap().alias("l")]) + .unwrap(); + let ran = block_on(df.count()); + let got = obs.get_timeout(Duration::from_secs(10)); + let e = ran + .err() + .or(got.err()) + .expect("the limit is exceeded by the retained states plus their merge"); + assert_eq!(e.error_class().name(), "UNABLE_TO_ACQUIRE_MEMORY", "{e}"); +} + +/// One observation whose metrics fail to finalize does not take the others' with it, and the +/// failure reaches the Connect sink instead of leaving it silently unset. +#[test] +fn a_failed_observation_reaches_the_sink_beside_the_good_ones() { + let s = Session::new().unwrap(); + s.conf().set("spark.sql.ansi.enabled", "true").unwrap(); + let sink = krishiv_session::ObservedSink::new(); + let df = base(&s) + .observe_named("good", [count(lit(1)).alias("n")]) + .observe_named( + "bad", + [expr("count(1) / (count(1) - 3)").unwrap().alias("r")], + ) + .with_observed_sink(&sink); + block_on(df.count()).unwrap(); + let all = sink.take().expect("the action finished"); + assert_eq!(metrics(all["good"].as_ref().unwrap())["n"], 3); + let e = all["bad"].as_ref().expect_err("the metric's own error"); + assert_eq!(e.error_class().name(), "DIVIDE_BY_ZERO", "{e}"); +} + +/// An observed frame registered as a temporary view completes its `Observation` when the view is +/// read through `Session::table` or SQL. +#[test] +fn a_temp_view_of_an_observed_frame_completes_the_observation() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(df.create_or_replace_temp_view("v")).unwrap(); + assert_eq!(block_on(s.table("v").count()).unwrap(), 3); + assert_eq!( + metrics(&obs.get_timeout(Duration::from_secs(10)).unwrap().unwrap())["n"], + 3 + ); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(df.create_or_replace_temp_view("w")).unwrap(); + block_on(async { s.sql("SELECT * FROM w").await.unwrap().count().await }).unwrap(); + assert_eq!( + metrics(&obs.get_timeout(Duration::from_secs(10)).unwrap().unwrap())["n"], + 3 + ); +} From 61891b720005dd6202d4399f4a5eaa639c658863 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 13:06:40 +0000 Subject: [PATCH 11/15] session: a view keeps only its Observation handles, not a script result's lease; refresh python public-api snapshot Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-session/src/dataframe.rs | 2 +- crates/krishiv-session/src/dataframe/observe.rs | 10 ++++++++++ python/public-api.txt | 2 +- 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index bd700f4c7..54510a2c9 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -1818,7 +1818,7 @@ impl DataFrame { self.session .state .catalog - .set_view_keeps(name, self.keep.clone()); + .set_view_keeps(name, self.observation_handles()); self.session.state.cache.forget_table(name); Ok(()) } diff --git a/crates/krishiv-session/src/dataframe/observe.rs b/crates/krishiv-session/src/dataframe/observe.rs index 11bb38ec9..c65be0c26 100644 --- a/crates/krishiv-session/src/dataframe/observe.rs +++ b/crates/krishiv-session/src/dataframe/observe.rs @@ -76,6 +76,16 @@ impl DataFrame { out } + /// The `Observation` handles this frame keeps (not its other keep-alive handles, such as the + /// lease of a script result, which a view must not hold). + pub(super) fn observation_handles(&self) -> Vec> { + self.keep + .iter() + .filter(|k| Arc::clone(k).downcast::().is_ok()) + .cloned() + .collect() + } + /// The observations and sinks this frame carries, and those of the temporary views it reads /// (`through_views`), for the query's env. pub(super) fn observers_through( diff --git a/python/public-api.txt b/python/public-api.txt index f4109a0a0..7e7d67923 100644 --- a/python/public-api.txt +++ b/python/public-api.txt @@ -84,7 +84,7 @@ DataFrame.def write_table(self, insert_into: bool, table: str, format: str | Non DataFrame.def write_to_table(self, table: str, mode: str, provider: str | None=..., options: list[tuple[str, str]] | None=..., properties: list[tuple[str, str]] | None=..., partitioned_by: list[str] | None=..., cluster_by: list[str] | None=..., condition: str | None=...) -> None GroupedData.def agg(self, cols: list[Column]) -> DataFrame Observation.@property def name(self) -> str -Observation.def __init__(self, name: str | None=...) -> None +Observation.def __new__(cls, name: str | None=...) -> Observation Observation.def get(self) -> ArrowStream Session.@staticmethod def is_dataframe_table(name: str) -> bool Session.def __new__(cls, config: list[tuple[str, str]] | None=...) -> Session From 71e34a5d47bd3fa2b1f5c7985be74d1cd9c60fb9 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 19:04:52 +0000 Subject: [PATCH 12/15] exec: observe partials merge in (partition, lane) order; refuse a Python UDF in an observed metric (P3-T084 review) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/aggregate/observe.rs | 6 ++- crates/krishiv-exec/src/context.rs | 51 +++++++++++++++++-- .../krishiv-exec/src/physical/instantiate.rs | 17 +++++++ .../tests/observe_python_udf.rs | 38 ++++++++++++++ 4 files changed, 107 insertions(+), 5 deletions(-) create mode 100644 crates/krishiv-session/tests/observe_python_udf.rs diff --git a/crates/krishiv-exec/src/aggregate/observe.rs b/crates/krishiv-exec/src/aggregate/observe.rs index 8b056c376..55756b6f2 100644 --- a/crates/krishiv-exec/src/aggregate/observe.rs +++ b/crates/krishiv-exec/src/aggregate/observe.rs @@ -157,6 +157,8 @@ struct Local { query: Arc, finalizer: Arc, partition: u32, + /// The driver's lane among the logical partitions' (0 when it is not one). + lane: usize, sent: bool, /// The table's bytes, charged to the driver's consumer. reservation: Reservation, @@ -174,7 +176,8 @@ impl Local { let states = self.table.state_batch(&[0])?; // The states move to the query's own charge: release the driver's first. self.reservation.shrink(self.reservation.size()); - self.query.observe_partial(&self.finalizer, states) + self.query + .observe_partial(&self.finalizer, (self.partition, self.lane), states) } } @@ -227,6 +230,7 @@ impl Operator for CollectMetrics { query: Arc::clone(ctx.query()), finalizer: Arc::clone(&self.finalizer), partition: ctx.partition(), + lane: ctx.lane().map_or(0, |(lane, _)| lane), sent: false, reservation: ctx.consumer().reservation(), claimed: None, diff --git a/crates/krishiv-exec/src/context.rs b/crates/krishiv-exec/src/context.rs index b37c9fff4..286d54b60 100644 --- a/crates/krishiv-exec/src/context.rs +++ b/crates/krishiv-exec/src/context.rs @@ -243,7 +243,48 @@ pub struct QueryContext { /// drivers. struct Observed { finalizer: Arc, + /// The drivers' states, kept in `(partition, lane)` order whatever order they arrived in, so + /// an order-sensitive metric (`collect_list`, `first`, a float `sum`) is the same every run. partials: Vec, + /// The `(partition, lane)` of each of `partials`, ascending. + order: Vec<(u32, usize)>, +} + +/// Puts `item` at the place `key` sorts to in `keys`, keeping `items` parallel to them; equal keys +/// keep their arrival order. +fn insert_in_order( + keys: &mut Vec<(u32, usize)>, + items: &mut Vec, + key: (u32, usize), + item: T, +) { + let at = keys.partition_point(|k| *k <= key); + keys.insert(at, key); + items.insert(at, item); +} + +#[cfg(test)] +mod observe_order_tests { + use super::insert_in_order; + + #[test] + fn partials_merge_in_partition_and_lane_order_whatever_the_arrival_order() { + let sources = [(0, 0), (0, 1), (1, 0), (2, 0), (2, 1)]; + for rot in 0..sources.len() { + for rev in [false, true] { + let mut arrival = sources.to_vec(); + arrival.rotate_left(rot); + if rev { + arrival.reverse(); + } + let (mut keys, mut items) = (Vec::new(), Vec::new()); + for k in arrival { + insert_in_order(&mut keys, &mut items, k, k); + } + assert_eq!(items, sources); + } + } + } } impl fmt::Debug for QueryContext { @@ -298,6 +339,7 @@ impl QueryContext { pub(crate) fn observe_partial( &self, finalizer: &Arc, + source: (u32, usize), states: arrow_array::RecordBatch, ) -> Result<()> { let bytes = states.get_array_memory_size(); @@ -306,14 +348,15 @@ impl QueryContext { let reservation = held.get_or_insert_with(|| self.observed_consumer().reservation()); reservation.try_grow(bytes)?; } - lock(&self.observed) + let mut all = lock(&self.observed); + let o = all .entry((finalizer.name().to_owned(), finalizer.instance())) .or_insert_with(|| Observed { finalizer: Arc::clone(finalizer), partials: Vec::new(), - }) - .partials - .push(states); + order: Vec::new(), + }); + insert_in_order(&mut o.order, &mut o.partials, source, states); Ok(()) } diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index 1c6e7b7b6..d7a618da6 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -1759,6 +1759,14 @@ pub(crate) fn collect_metrics_op( let s = p.schema(); (Some(p), s) } + // A registered Python UDF in a metric's argument plans a `PythonEval` stage under the + // projection; the observing operator folds rows in place and does not run stages. + other if has_python_eval(other) => { + return Err(krishiv_common::error::KrishivUnsupportedFeature { + feature: "a Python UDF in an observed metric (Dataset.observe)".to_owned(), + } + .into()); + } _ => return Err(bad()), }; let partial = AggregatePlan { @@ -1787,6 +1795,15 @@ pub(crate) fn collect_metrics_op( .map(|op| op as Arc) } +/// Whether `node` or a node below it is a `PythonEval`. +fn has_python_eval(node: &PhysNode) -> bool { + match node { + PhysNode::PythonEval { .. } => true, + PhysNode::Project { input, .. } => has_python_eval(input), + _ => false, + } +} + pub(crate) fn replicate_op( schema: &SchemaRef, layout: &[AttrId], diff --git a/crates/krishiv-session/tests/observe_python_udf.rs b/crates/krishiv-session/tests/observe_python_udf.rs new file mode 100644 index 000000000..219386fee --- /dev/null +++ b/crates/krishiv-session/tests/observe_python_udf.rs @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +//! A registered Python UDF in an observed metric's argument (P3-T084): the observing operator +//! does not run `PythonEval` stages, so it is refused with a clear error, not `INTERNAL_ERROR`. + +use krishiv_plan::PythonEvalType; +use krishiv_session::functions::{call_udf, col, sum}; +use krishiv_session::{Observation, PythonUdfDef, Session}; +use krishiv_types::DataType; + +#[tokio::test] +async fn a_python_udf_in_a_metric_is_unsupported_not_internal() { + let s = Session::new().unwrap(); + s.register_python_udf(PythonUdfDef { + name: "py_plus".to_owned(), + eval_type: PythonEvalType::ArrowBatched, + return_type: DataType::Long, + deterministic: true, + thread_safe: true, + }); + let obs = Observation::new(); + let df = s + .sql("SELECT id FROM range(5)") + .await + .unwrap() + .observe(&obs, [sum(call_udf("py_plus", [col("id")])).alias("s")]) + .unwrap(); + let err = match df.count().await { + Ok(n) => panic!("expected an error, counted {n}"), + Err(e) => e, + }; + let text = err.to_string(); + assert!(text.contains("Python UDF"), "{text}"); + assert_ne!( + format!("{:?}", err.error_class()), + "InternalError", + "{text}" + ); +} From f2c8f6e908fec92a4658698cf7677d36d8f40c4f Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 20:25:52 +0000 Subject: [PATCH 13/15] exec: move the CollectMetrics builder to instantiate/observe.rs (file-size gate, P3-T084) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/physical/instantiate.rs | 104 +---------------- .../src/physical/instantiate/observe.rs | 106 ++++++++++++++++++ 2 files changed, 108 insertions(+), 102 deletions(-) create mode 100644 crates/krishiv-exec/src/physical/instantiate/observe.rs diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index d7a618da6..2c07d0221 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -91,8 +91,10 @@ mod order; mod python_eval; pub(crate) use detached::{Detached, LoopRun}; mod join; +mod observe; use distributed::planned_schema_holds; pub(crate) use join::hash_join; +pub(crate) use observe::collect_metrics_op; use order::{range_ordered_evals, reads_range_in_order}; /// Creates the spill scope of one instantiation: the query's own spill @@ -1702,108 +1704,6 @@ pub(crate) fn expand_op( Ok(Expand::try_new(projects)?) } -/// The operator of a `CollectMetrics` node over rows of `schema`: the -/// shape of `measure` is what `planner/observe.rs` makes. -pub(crate) fn collect_metrics_op( - name: &str, - instance: u32, - limit: Option, - measure: &PhysNode, - schema: &SchemaRef, - v: &Values<'_>, -) -> Result> { - let bad = || { - internal(format!( - "the metrics of observation {name} have an unexpected plan" - )) - }; - let PhysNode::Project { - input: aggregate, - layout: post_layout, - items: post_items, - nullable: post_nullable, - subqueries, - } = measure - else { - return Err(bad()); - }; - let PhysNode::Aggregate { - input: args, - aggregate: plan, - } = aggregate.as_ref() - else { - return Err(bad()); - }; - if !subqueries.is_empty() { - return Err(krishiv_common::error::KrishivUnsupportedFeature { - feature: "a subquery in an observed metric (Dataset.observe)".to_owned(), - } - .into()); - } - if !plan.group_by.is_empty() - || plan.output != AggregateOutput::Final - || plan.distinct.iter().any(|d| *d) - { - return Err(bad()); - } - let (pre, pre_schema) = match args.as_ref() { - PhysNode::OneRow => (None, Arc::clone(schema)), - PhysNode::Project { - input: leaf, - layout, - items, - nullable, - subqueries, - } if matches!(leaf.as_ref(), PhysNode::OneRow) && subqueries.is_empty() => { - let p = project_op(schema, layout, items, nullable.as_deref(), v)?; - let s = p.schema(); - (Some(p), s) - } - // A registered Python UDF in a metric's argument plans a `PythonEval` stage under the - // projection; the observing operator folds rows in place and does not run stages. - other if has_python_eval(other) => { - return Err(krishiv_common::error::KrishivUnsupportedFeature { - feature: "a Python UDF in an observed metric (Dataset.observe)".to_owned(), - } - .into()); - } - _ => return Err(bad()), - }; - let partial = AggregatePlan { - output: AggregateOutput::Partial, - ..plan.clone() - }; - let merge = final_of(&partial)?; - CollectMetrics::try_new( - name, - instance, - limit, - Arc::clone(schema), - pre, - aggregate_spec(&pre_schema, &partial, v.env, None), - |states| Ok(aggregate_spec(states, &merge, v.env, None)), - |results| { - project_op( - results, - post_layout, - post_items, - post_nullable.as_deref(), - v, - ) - }, - ) - .map(|op| op as Arc) -} - -/// Whether `node` or a node below it is a `PythonEval`. -fn has_python_eval(node: &PhysNode) -> bool { - match node { - PhysNode::PythonEval { .. } => true, - PhysNode::Project { input, .. } => has_python_eval(input), - _ => false, - } -} - pub(crate) fn replicate_op( schema: &SchemaRef, layout: &[AttrId], diff --git a/crates/krishiv-exec/src/physical/instantiate/observe.rs b/crates/krishiv-exec/src/physical/instantiate/observe.rs new file mode 100644 index 000000000..3de3489db --- /dev/null +++ b/crates/krishiv-exec/src/physical/instantiate/observe.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The operator of a `CollectMetrics` node (P3-T084 `Dataset.observe`). + +use super::*; + +/// The operator of a `CollectMetrics` node over rows of `schema`: the +/// shape of `measure` is what `planner/observe.rs` makes. +pub(crate) fn collect_metrics_op( + name: &str, + instance: u32, + limit: Option, + measure: &PhysNode, + schema: &SchemaRef, + v: &Values<'_>, +) -> Result> { + let bad = || { + internal(format!( + "the metrics of observation {name} have an unexpected plan" + )) + }; + let PhysNode::Project { + input: aggregate, + layout: post_layout, + items: post_items, + nullable: post_nullable, + subqueries, + } = measure + else { + return Err(bad()); + }; + let PhysNode::Aggregate { + input: args, + aggregate: plan, + } = aggregate.as_ref() + else { + return Err(bad()); + }; + if !subqueries.is_empty() { + return Err(krishiv_common::error::KrishivUnsupportedFeature { + feature: "a subquery in an observed metric (Dataset.observe)".to_owned(), + } + .into()); + } + if !plan.group_by.is_empty() + || plan.output != AggregateOutput::Final + || plan.distinct.iter().any(|d| *d) + { + return Err(bad()); + } + let (pre, pre_schema) = match args.as_ref() { + PhysNode::OneRow => (None, Arc::clone(schema)), + PhysNode::Project { + input: leaf, + layout, + items, + nullable, + subqueries, + } if matches!(leaf.as_ref(), PhysNode::OneRow) && subqueries.is_empty() => { + let p = project_op(schema, layout, items, nullable.as_deref(), v)?; + let s = p.schema(); + (Some(p), s) + } + // A registered Python UDF in a metric's argument plans a `PythonEval` stage under the + // projection; the observing operator folds rows in place and does not run stages. + other if has_python_eval(other) => { + return Err(krishiv_common::error::KrishivUnsupportedFeature { + feature: "a Python UDF in an observed metric (Dataset.observe)".to_owned(), + } + .into()); + } + _ => return Err(bad()), + }; + let partial = AggregatePlan { + output: AggregateOutput::Partial, + ..plan.clone() + }; + let merge = final_of(&partial)?; + CollectMetrics::try_new( + name, + instance, + limit, + Arc::clone(schema), + pre, + aggregate_spec(&pre_schema, &partial, v.env, None), + |states| Ok(aggregate_spec(states, &merge, v.env, None)), + |results| { + project_op( + results, + post_layout, + post_items, + post_nullable.as_deref(), + v, + ) + }, + ) + .map(|op| op as Arc) +} + +/// Whether `node` or a node below it is a `PythonEval`. +fn has_python_eval(node: &PhysNode) -> bool { + match node { + PhysNode::PythonEval { .. } => true, + PhysNode::Project { input, .. } => has_python_eval(input), + _ => false, + } +} From 61af3039c4c30e7cc820d74bcd35c6e9a00b3b2b Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 21:09:12 +0000 Subject: [PATCH 14/15] session: a view over a view of an observed frame keeps the inherited observation handles (P3-T084 review) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-session/src/dataframe.rs | 2 +- .../krishiv-session/src/dataframe/observe.rs | 19 +++++++++++++++++++ crates/krishiv-session/tests/observe.rs | 16 ++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/crates/krishiv-session/src/dataframe.rs b/crates/krishiv-session/src/dataframe.rs index f6a11953f..0f05bacff 100644 --- a/crates/krishiv-session/src/dataframe.rs +++ b/crates/krishiv-session/src/dataframe.rs @@ -1818,7 +1818,7 @@ impl DataFrame { self.session .state .catalog - .set_view_keeps(name, self.observation_handles()); + .set_view_keeps(name, self.observation_handles_with(&analysis.env.observers)); self.session.state.cache.forget_table(name); Ok(()) } diff --git a/crates/krishiv-session/src/dataframe/observe.rs b/crates/krishiv-session/src/dataframe/observe.rs index c65be0c26..fb9db3f97 100644 --- a/crates/krishiv-session/src/dataframe/observe.rs +++ b/crates/krishiv-session/src/dataframe/observe.rs @@ -86,6 +86,25 @@ impl DataFrame { .collect() } + /// As [`Self::observation_handles`], plus the observations the analysis inherited through the + /// temporary views the plan reads (`inherited`), so a view over a view completes them too. + pub(super) fn observation_handles_with( + &self, + inherited: &Observers, + ) -> Vec> { + let mut out = self.observation_handles(); + for shared in &inherited.observations { + if !self.keep.iter().any(|k| { + Arc::clone(k) + .downcast::() + .is_ok_and(|s| Arc::ptr_eq(&s, shared)) + }) { + out.push(Arc::clone(shared) as _); + } + } + out + } + /// The observations and sinks this frame carries, and those of the temporary views it reads /// (`through_views`), for the query's env. pub(super) fn observers_through( diff --git a/crates/krishiv-session/tests/observe.rs b/crates/krishiv-session/tests/observe.rs index 26f1fe5b0..bb9cbcb2e 100644 --- a/crates/krishiv-session/tests/observe.rs +++ b/crates/krishiv-session/tests/observe.rs @@ -464,3 +464,19 @@ fn a_temp_view_of_an_observed_frame_completes_the_observation() { 3 ); } + +/// A view of a view of an observed frame completes the observation too, though the first view was +/// never executed. +#[test] +fn a_view_over_a_view_of_an_observed_frame_completes_the_observation() { + let s = Session::new().unwrap(); + let obs = Observation::new(); + let df = base(&s).observe(&obs, [count(lit(1)).alias("n")]).unwrap(); + block_on(df.create_or_replace_temp_view("v1")).unwrap(); + block_on(s.table("v1").create_or_replace_temp_view("v2")).unwrap(); + assert_eq!(block_on(s.table("v2").count()).unwrap(), 3); + assert_eq!( + metrics(&obs.get_timeout(Duration::from_secs(10)).unwrap().unwrap())["n"], + 3 + ); +} From aaa6ca222afc177290fa78c3a13879c804c7b4a2 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 22:28:20 +0000 Subject: [PATCH 15/15] plan, exec, scheduler: observe determinism stops at aggregates; map-task observations reach the query (P3-T084 review) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/context.rs | 18 +++++ crates/krishiv-plan/src/analyzer/observe.rs | 12 ++- crates/krishiv-plan/src/tests/observe.rs | 8 +- crates/krishiv-scheduler/src/stages.rs | 2 + crates/krishiv-scheduler/tests/tpch_stages.rs | 77 ++++++++++++++++++- 5 files changed, 114 insertions(+), 3 deletions(-) diff --git a/crates/krishiv-exec/src/context.rs b/crates/krishiv-exec/src/context.rs index 286d54b60..8e55815cd 100644 --- a/crates/krishiv-exec/src/context.rs +++ b/crates/krishiv-exec/src/context.rs @@ -360,6 +360,24 @@ impl QueryContext { Ok(()) } + /// Moves this context's observation partials to `parent`, for a task of a staged run whose + /// own context ends before the query's (the parent's `observed` then reports them). The states + /// are charged to `parent`'s pool and released from this one's. + /// + /// # Errors + /// `UNABLE_TO_ACQUIRE_MEMORY` when the parent's pool cannot hold them. + #[doc(hidden)] + pub fn forward_observed_to(&self, parent: &Self) -> Result<()> { + let taken = std::mem::take(&mut *lock(&self.observed)); + for o in taken.into_values() { + for (source, states) in o.order.into_iter().zip(o.partials) { + parent.observe_partial(&o.finalizer, source, states)?; + } + } + *lock(&self.observed_memory) = None; + Ok(()) + } + fn observed_consumer(&self) -> Arc { ConsumerHandle::register(Arc::clone(&self.pool), "observed-metrics", None) } diff --git a/crates/krishiv-plan/src/analyzer/observe.rs b/crates/krishiv-plan/src/analyzer/observe.rs index a29cb7232..501364027 100644 --- a/crates/krishiv-plan/src/analyzer/observe.rs +++ b/crates/krishiv-plan/src/analyzer/observe.rs @@ -67,7 +67,7 @@ fn violation(metric: &Arc) -> Option { return Some(Violation::Attribute); } ExprNode::Resolved(ResolvedExpr::Alias { .. }) => {} - _ if !seen && !crate::raises::is_deterministic(&e) => { + _ if !seen && is_nondeterministic_call(&e) => { return Some(Violation::NonDeterministic); } _ => {} @@ -79,6 +79,16 @@ fn violation(metric: &Arc) -> Option { None } +/// Whether `e` itself, not an expression below it, is a nondeterministic (or unresolved) call: the +/// walk in [`violation`] visits the children, and stops counting once inside an aggregate. +fn is_nondeterministic_call(e: &Expr) -> bool { + match &e.node { + ExprNode::Function(_) => true, + ExprNode::Resolved(ResolvedExpr::Function(f)) => !f.deterministic, + _ => false, + } +} + /// Whether `e` holds a subquery expression. fn has_subquery(e: &Arc) -> bool { let mut stack = vec![Arc::clone(e)]; diff --git a/crates/krishiv-plan/src/tests/observe.rs b/crates/krishiv-plan/src/tests/observe.rs index 28241ef49..03c4a8dce 100644 --- a/crates/krishiv-plan/src/tests/observe.rs +++ b/crates/krishiv-plan/src/tests/observe.rs @@ -47,6 +47,12 @@ fn allowed_metrics_pass_rows_through() { int(42), f("sum", vec![f("rand", vec![])]).alias("r"), f("max", vec![c("p_partkey")]), + // A nondeterministic call inside an aggregate stays allowed under arithmetic. + crate::Expr::builtin( + crate::Builtin::Add, + vec![f("sum", vec![f("rand", vec![])]), int(1)], + ) + .alias("r2"), ], ), &ctx(), @@ -56,7 +62,7 @@ fn allowed_metrics_pass_rows_through() { let PlanNode::CollectMetrics { metrics, input, .. } = &root.node else { panic!("root is {root}"); }; - assert_eq!(metrics.len(), 5); + assert_eq!(metrics.len(), 6); assert_eq!( root.resolved_output().map(|o| o.len()), input.resolved_output().map(|o| o.len()) diff --git a/crates/krishiv-scheduler/src/stages.rs b/crates/krishiv-scheduler/src/stages.rs index 41bfa9ed7..02c36397a 100644 --- a/crates/krishiv-scheduler/src/stages.rs +++ b/crates/krishiv-scheduler/src/stages.rs @@ -387,6 +387,8 @@ impl StageRunner { self.record(&task.metrics()); outcome?; task.error().map_or(Ok(()), Err)?; + // Its `CollectMetrics` partials belong to the query, which outlives the task. + task.forward_observed_to(query)?; Ok(produced(&stage.plan, values.as_ref())) } diff --git a/crates/krishiv-scheduler/tests/tpch_stages.rs b/crates/krishiv-scheduler/tests/tpch_stages.rs index 8962afa54..09aa9bd62 100644 --- a/crates/krishiv-scheduler/tests/tpch_stages.rs +++ b/crates/krishiv-scheduler/tests/tpch_stages.rs @@ -39,7 +39,7 @@ use krishiv_plan::analyzer::testing::{block_on, context}; use krishiv_plan::analyzer::{ AnalysisContext, Analyzer as _, DefaultAnalyzer, MemoryCatalog, ResolutionConf, TableColumn, }; -use krishiv_plan::{Plan, SourceId}; +use krishiv_plan::{Expr, LogicalPlan, Plan, PlanNode, SourceId}; use krishiv_scheduler::{StageRunner, StageSplits, TableDealer, fragment}; use object_store::ObjectStore; use object_store::local::LocalFileSystem; @@ -884,3 +884,78 @@ fn a_join_and_aggregate_run_in_pools_between_spilling_and_fitting() { } } } + +/// Rows of `lineitem` grouped by `l_returnflag` (a hash exchange), observed +/// below the exchange as `m`. +fn observed_group_by(tpch: &Tpch) -> PhysicalPlan { + let observed = LogicalPlan::new(PlanNode::CollectMetrics { + input: Arc::new(LogicalPlan::table(["lineitem"])), + name: "m".to_owned(), + metrics: vec![Arc::new( + Expr::function("count", vec![Expr::column(["l_orderkey"])]).alias("n"), + )], + }); + let logical = observed.aggregate( + vec![Expr::column(["l_returnflag"])], + vec![ + Expr::column(["l_returnflag"]), + Expr::function("count", vec![Expr::column(["l_orderkey"])]).alias("c"), + ], + ); + let plan = Plan::try_new(logical).unwrap(); + let plan: Plan = block_on(DefaultAnalyzer.analyze(&plan, &tpch.ctx)) + .unwrap() + .plan; + let cx = PlanningContext::new(Arc::clone(&tpch.tables) as _, Arc::clone(&tpch.config)); + DefaultPlanner.plan(&plan, &cx).unwrap() +} + +#[test] +fn an_observation_below_an_exchange_reaches_the_query() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + let tpch = tpch(&rt); + let executor = TokioExecutor::new(&ExecutorConfig::default()).unwrap(); + let stages = StageRunner::new(Arc::clone(&executor) as _, 4).unwrap(); + let root: Arc = TreePool::new(&MemoryConfig::default()); + let worker = WorkerContext::new(Arc::clone(&tpch.tables) as _); + let plan = observed_group_by(&tpch); + let query = QueryContext::new( + QueryId(9), + WorkloadClass::Interactive, + ExecConfig::default(), + root.child("q9", None), + ); + let grouped: i64 = rt.block_on(async { + let mut stream = stages + .submit(plan.clone(), worker.clone(), Arc::clone(&query)) + .await + .unwrap(); + let mut total = 0; + while let Some(b) = stream.next().await { + let batch = b.unwrap().into_parts().0; + let counts = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + total += counts.iter().flatten().sum::(); + } + query.finished().await; + total + }); + assert!(grouped > 0); + let observed = query.observed().unwrap(); + let m = observed + .get("m") + .expect("the observation reached the query"); + let n = m + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(n, grouped); +}