%%% @doc Query Planner and Prepared Statement Execution
%%%
%%% Provides a Lisp-inspired tuple-based DSL for constructing and executing
%%% relational algebra queries. Queries are represented as nested Erlang tuples
%%% that can be validated, optimized, and compiled into iterator pipelines.
%%%
%%% == Query DSL Syntax ==
%%%
%%% The DSL uses nested tuples to represent relational operations:
%%%
%%% ```
%%% % Scan relation (base case)
%%% {scan, employees}
%%%
%%% % Reference relation (synonym for scan)
%%% {relation, employees}
%%%
%%% % Selection (filter)
%%% {select, {scan, employees}, fun(T) -> maps:get(age, T) > 30 end}
%%%
%%% % Projection
%%% {project, {scan, employees}, [name, age]}
%%%
%%% % Equijoin (natural join on common attribute)
%%% {join, {scan, employees}, {scan, departments}, dept_id}
%%%
%%% % Theta join (join with arbitrary predicate)
%%% {theta_join,
%%% {scan, employees},
%%% {scan, departments},
%%% fun(L, R) -> maps:get(dept_id, L) =:= maps:get(id, R) end}
%%%
%%% % Sort
%%% {sort, {scan, employees}, fun(A, B) -> maps:get(age, A) =< maps:get(age, B) end}
%%%
%%% % Take (limit)
%%% {take, {scan, employees}, 10}
%%%
%%% % Rename attribute
%%% {rename, old_name, new_name, {scan, employees}}
%%%
%%% % Composite query
%%% {project,
%%% {select,
%%% {join,
%%% {scan, employees},
%%% {scan, departments},
%%% dept_id},
%%% fun(T) -> maps:get(age, T) > 30 end},
%%% [name, dept_name]}
%%% ```
%%%
%%% == Usage ==
%%%
%%% ```
%%% % Prepare a query plan
%%% Plan = query_planner:prepare({select, {scan, employees},
%%% fun(T) -> maps:get(age, T) > 30 end}).
%%%
%%% % Execute and get iterator
%%% Iter = query_planner:execute(DB, Plan).
%%%
%%% % Collect all results
%%% Results = query_planner:collect(DB, Plan).
%%% ```
%%%
%%% @author Nekoma Team
%%% @copyright 2025
-module(query_planner).
-export([prepare/1,
execute/2,
collect/2,
explain/1]).
-include("operations.hrl").
%%% Types
-type query_plan() ::
{scan, atom()} % Scan relation
| {relation, atom()} % Reference relation
| {select, query_plan(), fun((map()) -> boolean())} % Filter
| {project, query_plan(), [atom()]} % Projection
| {join, query_plan(), query_plan(), atom()} % Equijoin
| {theta_join, query_plan(), query_plan(),
fun((map(), map()) -> boolean())} % Theta join
| {sort, query_plan(), fun((map(), map()) -> boolean())} % Sort
| {take, query_plan(), pos_integer()} % Limit
| {rename, atom(), atom(), query_plan()} % Rename attribute
| {materialize, query_plan()} % Materialize all
| {materialize, query_plan(), pos_integer()}. % Materialize N tuples
-export_type([query_plan/0]).
%%% Public API
%% @doc Prepare and validate a query plan.
%%
%% Takes a query plan tuple and validates its structure. Returns the
%% plan unchanged if valid, or throws an error if invalid.
%%
%% @param Plan The query plan to prepare
%% @returns The validated query plan
%% @throws {invalid_plan, Reason}
-spec prepare(query_plan()) -> query_plan().
prepare(Plan) ->
case validate_plan(Plan) of
ok -> Plan;
{error, Reason} -> error({invalid_plan, Reason})
end.
%% @doc Execute a prepared query plan against a database.
%%
%% Compiles the query plan into an iterator pipeline and returns
%% the iterator handle. The caller must consume the iterator using
%% operations:next_tuple/1 and close it with operations:close_iterator/1.
%%
%% @param DB The database state
%% @param Plan The prepared query plan
%% @returns Iterator process PID
-spec execute(term(), query_plan()) -> pid().
execute(DB, Plan) ->
compile_to_iterator(DB, Plan).
%% @doc Execute a query plan and collect all results.
%%
%% Convenience function that executes the plan and collects all
%% tuples into a list. Automatically closes the iterator.
%%
%% If the iterator times out (e.g., when sorting an infinite relation),
%% returns partial results with a warning.
%%
%% @param DB The database state
%% @param Plan The prepared query plan
%% @returns List of result tuples, or partial results on timeout
-spec collect(term(), query_plan()) -> [map()].
collect(DB, Plan) ->
Iter = execute(DB, Plan),
Results = operations:collect_all(Iter),
operations:close_iterator(Iter),
case Results of
{error, timeout, PartialResults} ->
io:format("~nWarning: Query timed out. Returning partial results (~b tuples).~n",
[length(PartialResults)]),
io:format("Tip: For infinite relations, use 'take' before operations like 'sort'.~n~n"),
PartialResults;
{error, Reason, PartialResults} ->
io:format("~nWarning: Query failed (~p). Returning partial results (~b tuples).~n~n",
[Reason, length(PartialResults)]),
PartialResults;
_ ->
Results
end.
%% @doc Explain a query plan in human-readable form.
%%
%% Converts the query plan to a readable string representation
%% showing the operation tree.
%%
%% @param Plan The query plan to explain
%% @returns String representation of the plan
-spec explain(query_plan()) -> string().
explain(Plan) ->
explain_plan(Plan, 0).
%%% Internal Functions
%% Validate query plan structure
-spec validate_plan(query_plan()) -> ok | {error, term()}.
validate_plan({scan, Name}) when is_atom(Name) ->
ok;
validate_plan({relation, Name}) when is_atom(Name) ->
ok;
validate_plan({select, SubPlan, Pred}) when is_function(Pred, 1) ->
validate_plan(SubPlan);
validate_plan({project, SubPlan, Attrs}) when is_list(Attrs) ->
case lists:all(fun is_atom/1, Attrs) of
true -> validate_plan(SubPlan);
false -> {error, {invalid_attributes, Attrs}}
end;
validate_plan({join, Left, Right, Attr}) when is_atom(Attr) ->
case {validate_plan(Left), validate_plan(Right)} of
{ok, ok} -> ok;
{{error, _} = Error, _} -> Error;
{_, {error, _} = Error} -> Error
end;
validate_plan({theta_join, Left, Right, Pred}) when is_function(Pred, 2) ->
case {validate_plan(Left), validate_plan(Right)} of
{ok, ok} -> ok;
{{error, _} = Error, _} -> Error;
{_, {error, _} = Error} -> Error
end;
validate_plan({sort, SubPlan, Comp}) when is_function(Comp, 2) ->
validate_plan(SubPlan);
validate_plan({take, SubPlan, N}) when is_integer(N), N > 0 ->
validate_plan(SubPlan);
validate_plan({rename, OldAttr, NewAttr, SubPlan}) when is_atom(OldAttr), is_atom(NewAttr) ->
validate_plan(SubPlan);
validate_plan({materialize, SubPlan}) ->
validate_plan(SubPlan);
validate_plan({materialize, SubPlan, N}) when is_integer(N), N > 0 ->
validate_plan(SubPlan);
validate_plan(Plan) ->
{error, {unknown_operator, Plan}}.
%% Compile query plan to iterator pipeline
-spec compile_to_iterator(term(), query_plan()) -> pid().
compile_to_iterator(DB, {scan, Name}) ->
operations:get_tuples_iterator(DB, Name);
compile_to_iterator(DB, {relation, Name}) ->
operations:get_tuples_iterator(DB, Name);
compile_to_iterator(DB, {select, SubPlan, Predicate}) ->
SubIter = compile_to_iterator(DB, SubPlan),
operations:select_iterator(SubIter, Predicate);
compile_to_iterator(DB, {project, SubPlan, Attributes}) ->
SubIter = compile_to_iterator(DB, SubPlan),
operations:project_iterator(SubIter, Attributes);
compile_to_iterator(DB, {join, LeftPlan, RightPlan, Attribute}) ->
LeftIter = compile_to_iterator(DB, LeftPlan),
RightIter = compile_to_iterator(DB, RightPlan),
operations:equijoin_iterator(LeftIter, RightIter, Attribute);
compile_to_iterator(DB, {theta_join, LeftPlan, RightPlan, Predicate}) ->
LeftIter = compile_to_iterator(DB, LeftPlan),
RightIter = compile_to_iterator(DB, RightPlan),
operations:theta_join_iterator(LeftIter, RightIter, Predicate);
compile_to_iterator(DB, {sort, SubPlan, Comparator}) ->
SubIter = compile_to_iterator(DB, SubPlan),
operations:sort_iterator(SubIter, Comparator);
compile_to_iterator(DB, {take, SubPlan, N}) ->
SubIter = compile_to_iterator(DB, SubPlan),
operations:take_iterator(SubIter, N);
compile_to_iterator(DB, {rename, OldAttr, NewAttr, SubPlan}) ->
SubIter = compile_to_iterator(DB, SubPlan),
rename_iterator(SubIter, OldAttr, NewAttr);
compile_to_iterator(DB, {materialize, SubPlan}) ->
% Materialize returns all tuples as a "static" iterator
SubIter = compile_to_iterator(DB, SubPlan),
Results = operations:collect_all(SubIter),
operations:close_iterator(SubIter),
static_iterator(Results);
compile_to_iterator(DB, {materialize, SubPlan, Limit}) ->
% Materialize with limit
SubIter = compile_to_iterator(DB, SubPlan),
LimitedIter = operations:take_iterator(SubIter, Limit),
Results = operations:collect_all(LimitedIter),
operations:close_iterator(LimitedIter),
static_iterator(Results).
%% Static iterator - serves pre-materialized tuples
-spec static_iterator([map()]) -> pid().
static_iterator(Tuples) ->
spawn(fun() -> static_loop(Tuples) end).
static_loop(Tuples) ->
receive
{next, Caller} ->
case Tuples of
[] ->
Caller ! done,
static_loop([]);
[Tuple | Rest] ->
Caller ! {tuple, Tuple},
static_loop(Rest)
end;
{close, Caller} ->
Caller ! ok
end.
%% Rename iterator - renames an attribute in each tuple
-spec rename_iterator(pid(), atom(), atom()) -> pid().
rename_iterator(SourceIter, OldAttr, NewAttr) ->
spawn(fun() -> rename_loop(SourceIter, OldAttr, NewAttr) end).
rename_loop(SourceIter, OldAttr, NewAttr) ->
receive
{next, Caller} ->
case operations:next_tuple(SourceIter) of
done ->
Caller ! done,
rename_loop(SourceIter, OldAttr, NewAttr);
{ok, Tuple} ->
RenamedTuple = case maps:is_key(OldAttr, Tuple) of
true ->
Value = maps:get(OldAttr, Tuple),
Tuple1 = maps:remove(OldAttr, Tuple),
maps:put(NewAttr, Value, Tuple1);
false ->
Tuple
end,
Caller ! {tuple, RenamedTuple},
rename_loop(SourceIter, OldAttr, NewAttr)
end;
{close, Caller} ->
operations:close_iterator(SourceIter),
Caller ! ok
end.
%% Pretty-print query plan
-spec explain_plan(query_plan(), non_neg_integer()) -> string().
explain_plan({scan, Name}, Indent) ->
indent(Indent) ++ "Scan: " ++ atom_to_list(Name) ++ "\n";
explain_plan({relation, Name}, Indent) ->
indent(Indent) ++ "Relation: " ++ atom_to_list(Name) ++ "\n";
explain_plan({select, SubPlan, _Pred}, Indent) ->
indent(Indent) ++ "Select (filter)\n" ++
explain_plan(SubPlan, Indent + 2);
explain_plan({project, SubPlan, Attrs}, Indent) ->
AttrStr = string:join([atom_to_list(A) || A <- Attrs], ", "),
indent(Indent) ++ "Project [" ++ AttrStr ++ "]\n" ++
explain_plan(SubPlan, Indent + 2);
explain_plan({join, LeftPlan, RightPlan, Attr}, Indent) ->
indent(Indent) ++ "Join on " ++ atom_to_list(Attr) ++ "\n" ++
explain_plan(LeftPlan, Indent + 2) ++
explain_plan(RightPlan, Indent + 2);
explain_plan({theta_join, LeftPlan, RightPlan, _Pred}, Indent) ->
indent(Indent) ++ "Theta Join\n" ++
explain_plan(LeftPlan, Indent + 2) ++
explain_plan(RightPlan, Indent + 2);
explain_plan({sort, SubPlan, _Comp}, Indent) ->
indent(Indent) ++ "Sort\n" ++
explain_plan(SubPlan, Indent + 2);
explain_plan({take, SubPlan, N}, Indent) ->
indent(Indent) ++ "Take " ++ integer_to_list(N) ++ "\n" ++
explain_plan(SubPlan, Indent + 2);
explain_plan({rename, OldAttr, NewAttr, SubPlan}, Indent) ->
indent(Indent) ++ "Rename " ++ atom_to_list(OldAttr) ++ " -> " ++ atom_to_list(NewAttr) ++ "\n" ++
explain_plan(SubPlan, Indent + 2);
explain_plan({materialize, SubPlan}, Indent) ->
indent(Indent) ++ "Materialize (all)\n" ++
explain_plan(SubPlan, Indent + 2);
explain_plan({materialize, SubPlan, N}, Indent) ->
indent(Indent) ++ "Materialize (" ++ integer_to_list(N) ++ " tuples)\n" ++
explain_plan(SubPlan, Indent + 2).
%% Helper: indent string
-spec indent(non_neg_integer()) -> string().
indent(N) -> lists:duplicate(N, $\s).
