From 6b98b15447841c8af81b7e5a47f19e775aaa06fb Mon Sep 17 00:00:00 2001 From: David Anthoff Date: Wed, 16 Sep 2026 22:47:14 -0700 Subject: [PATCH] Push the new LINQ operators down into SQL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Teaches the DuckDB backend about the operators added in QueryOperators 1.2, and fixes two things the new operators made untenable. Nullability. DuckDB hands back columns typed Union{Missing,T}, but Query's operators use DataValue for an absent value and never Missing, so the two backends disagreed the moment a null appeared — which every outer join guarantees. Rows are now converted on the way out, so piping a DuckDB result back through query operators, or comparing it against the in-memory backend, sees the same thing either way. The column interfaces still speak Missing, because that is what the TableTraits _using_missing protocol is defined in terms of and what table sinks expect, so DataFrame output is unchanged. Unsupported operations. build_sql used to fail with a bare "Unsupported query operation: ", which told the user nothing. Each node with no SQL equivalent now explains the obstacle and suggests a way forward, usually materializing the query first. This also covers @groupjoin, @mapmany, @summarize and the pivots, which were already unsupported. Terminal operators. count, any, all, first, element_at, min_by and max_by become part of the SQL rather than pulling every row into Julia. Anything not handled falls through to QueryableBackend's default, which materializes and runs the in-memory implementation, so correctness never depends on that list being complete. This is also what makes @count work at all against DuckDB: it previously had no Queryable method and failed outright. New SQL: LEFT/RIGHT/FULL OUTER JOIN, UNION ALL/UNION/EXCEPT/INTERSECT and their DISTINCT ON key-based forms, ORDER BY ALL for @order, ORDER BY random() for @shuffle, GROUP BY with COUNT(*) for @count_by, and QUALIFY over ROW_NUMBER() for @take_last and @drop_last. The last of these takes its row count from COUNT(*) OVER () rather than a repeated subquery, so the inner query's positional parameters stay in order. Two-input nodes now name their right-hand source after the node's position in the walked tree instead of a hardcoded source_tbl_2, so a query containing more than one join or set operation no longer has them collide. Requires QueryableBackend 1.1, which is not registered yet, so CI here stays red until it is. Co-Authored-By: Claude Opus 5 --- Project.toml | 4 +- README.md | 28 +++- src/QueryDuckDB.jl | 3 + src/execution.jl | 14 +- src/result.jl | 66 ++++++-- src/scalar.jl | 108 +++++++++++++ src/sql_generation.jl | 195 ++++++++++++++++++++---- src/unsupported.jl | 86 +++++++++++ test/runtests.jl | 1 + test/test_new_operators.jl | 303 +++++++++++++++++++++++++++++++++++++ 10 files changed, 764 insertions(+), 44 deletions(-) create mode 100644 src/scalar.jl create mode 100644 src/unsupported.jl create mode 100644 test/test_new_operators.jl diff --git a/Project.toml b/Project.toml index 79afc3a..cb59c19 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ authors = ["David Anthoff "] [deps] DBInterface = "a10d1c49-ce27-4219-8d33-6db1a4562965" +DataValues = "e7dc6d0d-1eca-5fa6-8ad6-5aecde8b7ea5" DuckDB = "d2f5444f-75bc-4fdf-ac35-56f514c445e1" IteratorInterfaceExtensions = "82899510-4779-5014-852e-03e436cf321d" QueryableBackend = "0898d9ac-042b-5d31-8dcc-959fb3365f19" @@ -30,12 +31,13 @@ ExcelFiles = {url = "https://github.com/queryverse/ExcelFiles.jl", rev = "main"} [compat] CSVFiles = "1.0.2" DBInterface = "2" +DataValues = "0.4.4, 0.5, 1" DuckDB = "1" ExcelFiles = "1" FeatherFiles = "0.8, 1" IteratorInterfaceExtensions = "1" ParquetFiles = "0.2, 1" -QueryableBackend = "1" +QueryableBackend = "1.1" TableTraits = "1" Tables = "1" julia = "1.12" diff --git a/README.md b/README.md index 9514a23..bbdf956 100644 --- a/README.md +++ b/README.md @@ -111,6 +111,23 @@ Physical Plan: | `@unique(_.col)`, `@unique({_.a, _.b})` | `SELECT DISTINCT ON (...)` | | `@groupby` | `GROUP BY`; use aggregations and `key(_)` in the following `@map` | | `@join` | `INNER JOIN` | +| `@left_join`, `@right_join`, `@full_join` | `LEFT`/`RIGHT`/`FULL OUTER JOIN` | +| `@concat`, `@union`, `@except`, `@intersect` | `UNION ALL`, `UNION`, `EXCEPT`, `INTERSECT` | +| `@union_by`, `@except_by`, `@intersect_by` | `DISTINCT ON` with `IN`/`NOT IN` | +| `@order`, `@order_descending` | `ORDER BY ALL` | +| `@shuffle()` | `ORDER BY random()` | +| `@take_last`, `@drop_last` | `QUALIFY ROW_NUMBER() OVER ()` against `COUNT(*) OVER ()` | +| `@count_by` | `GROUP BY` with `COUNT(*)` | +| `@count`, `@any`, `@all` | `COUNT(*)`, `EXISTS`, `NOT EXISTS` | +| `@first`, `@element_at` | `LIMIT`, `LIMIT ... OFFSET` | +| `@min_by`, `@max_by` | `ORDER BY ... LIMIT 1` | + +The second operand of a join or set operation must be a DuckDB source too, so +write `df1 |> @duckdb() |> @union(df2 |> @duckdb())`. + +Terminal operators that are not in the table above — `@last`, `@single`, +`@contains`, `@aggregate`, `@sequence_equal` — still work: they materialize the +query and run the in-memory implementation. Common Julia functions are translated to their SQL equivalents (`uppercase`, `lowercase`, `strip`, `replace`, `startswith`, `occursin`, @@ -120,8 +137,15 @@ the aggregations `sum`, `mean`, `minimum`, `maximum`, `length` and ## Known limitations -- `@groupjoin` and `@mapmany` are not supported and throw a - `TranslationError`. +- Operators with no SQL equivalent throw a `TranslationError` naming the + operator and suggesting a way forward, usually materializing the query + first: `@groupjoin`, `@mapmany`, `@chunk`, `@aggregate_by`, `@take_while`, + `@drop_while`, `@reverse`, `@index`, `@append`, `@prepend`, `@zip`, + `@of_type`, `@cast`, `@summarize`, `@pivot_longer` and `@pivot_wider`. +- Set operations return rows in whatever order DuckDB produces, while the + in-memory implementation preserves first-seen order. +- `@shuffle` accepts no `rng` argument here, because the shuffling is done by + DuckDB's own random number generator. - A `@groupby` must be followed by a `@map` with aggregations; the group elements cannot be materialized as arrays. Three-argument `@groupby` requires a plain column key. diff --git a/src/QueryDuckDB.jl b/src/QueryDuckDB.jl index 1a6d7c1..1793f44 100644 --- a/src/QueryDuckDB.jl +++ b/src/QueryDuckDB.jl @@ -8,15 +8,18 @@ import TableTraits import DuckDB import DBInterface import Tables +import DataValues export @duckdb, @duckdbplan, @queryplan include("source.jl") include("source_detection.jl") include("expr_translation.jl") +include("unsupported.jl") include("sql_generation.jl") include("execution.jl") include("result.jl") +include("scalar.jl") include("plan.jl") end # module QueryDuckDB diff --git a/src/execution.jl b/src/execution.jl index 86da29f..ea8dd63 100644 --- a/src/execution.jl +++ b/src/execution.jl @@ -36,8 +36,12 @@ end """ register_query_sources(con, query_tree) -Register every :table source in the query tree with the DuckDB connection: -the root source as `source_tbl`, and each join's inner source as `source_tbl_2`. +Register every :table source in the query tree with the DuckDB connection: the +root source as `source_tbl`, and the second input of every two-input node — +joins, set operations — under the name `inner_table_name` derives from its +position in the walked tree. sql_generation.jl walks the tree the same way, so +the two agree on which name refers to which source even when a query contains +several joins or set operations. """ function register_query_sources(con, query_tree::QueryableBackend.Queryable) source = QueryableBackend.get_source(query_tree) @@ -49,11 +53,11 @@ function register_query_sources(con, query_tree::QueryableBackend.Queryable) register_source_data(con, source.original_source, "source_tbl") end - for node in QueryableBackend.walk_tree(query_tree) - if node isa QueryableBackend.QueryableJoin + for (i, node) in enumerate(QueryableBackend.walk_tree(query_tree)) + if node isa QueryableBackend.QueryableBinary inner_source = node.inner if inner_source isa DuckDBQueryableSource && inner_source.source_type == :table - register_source_data(con, inner_source.original_source, "source_tbl_2") + register_source_data(con, inner_source.original_source, inner_table_name(i)) end end end diff --git a/src/result.jl b/src/result.jl index 013a783..75d0e68 100644 --- a/src/result.jl +++ b/src/result.jl @@ -13,27 +13,75 @@ struct DuckDBQueryResult columns::NamedTuple end -# --- Iteration protocol (required by Tables.jl row iteration path) --- +# --- Row iteration --- +# +# DuckDB hands back columns typed `Union{Missing,T}`, but Query's operators use +# DataValue for an absent value and never Missing. Rows are therefore converted +# on the way out, so that piping a DuckDB result back through query operators — +# or comparing it element-wise against the in-memory backend — sees the same +# thing either way. +# +# The *column* interfaces below still speak Missing, because that is what the +# TableTraits `_using_missing` protocol is defined in terms of and what table +# sinks such as DataFrame expect. -function Base.iterate(r::DuckDBQueryResult) - rows = Tables.rows(r.columns) - return iterate(rows) +_datavalue_fieldtype(::Type{S}) where {S} = + Missing <: S ? DataValues.DataValue{Base.nonmissingtype(S)} : S + +function _datavalue_row_type(::Type{NamedTuple{names,types}}) where {names,types} + field_types = Tuple{(_datavalue_fieldtype(eltype(t)) for t in types.parameters)...} + return NamedTuple{names,field_types} +end + +_as_field(::Type{DataValues.DataValue{S}}, v) where {S} = + v === missing ? DataValues.DataValue{S}() : DataValues.DataValue{S}(v) + +_as_field(::Type{S}, v) where {S} = v + +struct DuckDBRowIterator{T,C<:NamedTuple} + columns::C + nrows::Int +end + +function row_iterator(r::DuckDBQueryResult) + C = typeof(r.columns) + T = _datavalue_row_type(C) + nrows = length(r.columns) == 0 ? 0 : length(first(r.columns)) + return DuckDBRowIterator{T,C}(r.columns, nrows) +end + +Base.eltype(::Type{DuckDBRowIterator{T,C}}) where {T,C} = T + +Base.IteratorSize(::Type{<:DuckDBRowIterator}) = Base.HasLength() + +Base.length(it::DuckDBRowIterator) = it.nrows + +@generated function _build_row(::Type{T}, columns::NamedTuple{names,types}, i::Int) where {T,names,types} + fields = [:( _as_field($(fieldtype(T, n)), columns.$(names[n])[i]) ) for n in 1:length(names)] + return :( T(($(fields...),)) ) end -function Base.iterate(r::DuckDBQueryResult, state) - rows = Tables.rows(r.columns) - return iterate(rows, state) +function Base.iterate(it::DuckDBRowIterator{T,C}, i::Int=1) where {T,C} + i > it.nrows && return nothing + return _build_row(T, it.columns, i), i + 1 end +# --- Iteration protocol --- + +Base.iterate(r::DuckDBQueryResult) = iterate(row_iterator(r)) + +Base.iterate(r::DuckDBQueryResult, state) = iterate(row_iterator(r), state) + Base.length(r::DuckDBQueryResult) = length(first(r.columns)) -Base.eltype(r::DuckDBQueryResult) = eltype(Tables.rows(r.columns)) + +Base.eltype(r::DuckDBQueryResult) = eltype(row_iterator(r)) # --- IteratorInterfaceExtensions --- IteratorInterfaceExtensions.isiterable(::DuckDBQueryResult) = true function IteratorInterfaceExtensions.getiterator(r::DuckDBQueryResult) - return Tables.rows(r.columns) + return row_iterator(r) end # --- TableTraits --- diff --git a/src/scalar.jl b/src/scalar.jl new file mode 100644 index 0000000..0f1ea98 --- /dev/null +++ b/src/scalar.jl @@ -0,0 +1,108 @@ +""" +Push-down for terminal operators — the ones that return a value rather than +another query. + +QueryableBackend hands every terminal operator to `execute_scalar`, which +dispatches on the root source type. Specialising it here lets a terminal +operator become part of the SQL instead of pulling every row into Julia first. +Anything not handled falls through to QueryableBackend's default, which +materializes and runs the in-memory implementation — so correctness never +depends on the list below being complete. +""" + +function QueryableBackend._execute_scalar(::DuckDBQueryableSource, q::QueryableBackend.QueryableScalar) + sql = _scalar_sql(q) + sql === nothing && return QueryableBackend._execute_scalar_fallback(q) + return _run_scalar(q, sql) +end + +# Returns the SQL for a terminal operator, or `nothing` when it has no +# translation and should fall back to the in-memory path. +function _scalar_sql(q::QueryableBackend.QueryableScalar) + params = Any[] + inner = generate_sql(q.source) + append!(params, inner.params) + + body = "($(inner.sql))" + + if q.op == :count + if isempty(q.args) + return SQLQuery("SELECT COUNT(*) FROM $body AS scalar_subq", params) + end + predicate = translate_filter_expr(q.args[2], params) + return SQLQuery("SELECT COUNT(*) FROM $body AS scalar_subq WHERE $predicate", params) + + elseif q.op == :any + if isempty(q.args) + return SQLQuery("SELECT EXISTS (SELECT 1 FROM $body AS scalar_subq)", params) + end + predicate = translate_filter_expr(q.args[2], params) + return SQLQuery("SELECT EXISTS (SELECT 1 FROM $body AS scalar_subq WHERE $predicate)", params) + + elseif q.op == :all + predicate = translate_filter_expr(q.args[2], params) + # True when no row violates the predicate. A NULL result from the + # predicate is not a violation, matching the in-memory behaviour of + # treating only an explicit `false` as a failure. + return SQLQuery("SELECT NOT EXISTS (SELECT 1 FROM $body AS scalar_subq WHERE NOT COALESCE($predicate, TRUE))", params) + + elseif q.op == :first + if isempty(q.args) + return SQLQuery("SELECT * FROM $body AS scalar_subq LIMIT 1", params) + end + predicate = translate_filter_expr(q.args[2], params) + return SQLQuery("SELECT * FROM $body AS scalar_subq WHERE $predicate LIMIT 1", params) + + elseif q.op == :element_at + n = q.args[1] + n < 1 && error("element_at was called with index $n; the index must be at least 1.") + return SQLQuery("SELECT * FROM $body AS scalar_subq LIMIT 1 OFFSET $(n - 1)", params) + + elseif q.op == :min_by || q.op == :max_by + key = translate_orderby_expr(q.args[2], params) + direction = q.op == :min_by ? "ASC" : "DESC" + return SQLQuery("SELECT * FROM $body AS scalar_subq ORDER BY $key $direction LIMIT 1", params) + end + + return nothing +end + +# Runs the scalar SQL and unwraps the result the way the corresponding +# in-memory operator would: a bare value for the aggregates, a row for the +# element-returning operators. +function _run_scalar(q::QueryableBackend.QueryableScalar, sql::SQLQuery) + db = DuckDB.DB() + con = DBInterface.connect(db) + + try + register_query_sources(con, q.source) + + result = DBInterface.execute(con, sql.sql, sql.params) + cols = Tables.columns(result) + col_names = Tuple(Tables.columnnames(cols)) + columns_nt = NamedTuple{col_names}(Tuple(collect(Tables.getcolumn(cols, n)) for n in col_names)) + + nrows = length(col_names) == 0 ? 0 : length(first(columns_nt)) + + if q.op in (:count, :any, :all) + return only(first(columns_nt)) + end + + # Element-returning operators: no row means the sequence was empty or + # nothing matched, which is an error for every one of them. + if nrows == 0 + if q.op == :element_at + error("element_at was called with index $(q.args[1]) on a sequence with fewer elements.") + elseif isempty(q.args) + error("$(q.op) was called on a sequence with no elements.") + else + error("$(q.op) was called on a sequence with no element matching the predicate.") + end + end + + return first(row_iterator(DuckDBQueryResult(columns_nt))) + finally + DBInterface.close!(con) + DBInterface.close!(db) + end +end diff --git a/src/sql_generation.jl b/src/sql_generation.jl index 7c450ff..a47d80d 100644 --- a/src/sql_generation.jl +++ b/src/sql_generation.jl @@ -75,9 +75,20 @@ function build_sql(nodes::Vector{QueryableBackend.Queryable}, params::Vector{Any end elseif node isa QueryableBackend.QueryableOrderBy - col = translate_orderby_expr(node.keySelector_expr, params) direction = node.descending ? "DESC" : "ASC" - push!(orderby_clauses, "$col $direction") + if is_identity_lambda(node.keySelector_expr) + # @order()/@order_descending() sort by whole rows, which DuckDB + # spells ORDER BY ALL. + push!(orderby_clauses, "ALL $direction") + else + col = translate_orderby_expr(node.keySelector_expr, params) + push!(orderby_clauses, "$col $direction") + end + + elseif node isa QueryableBackend.QueryableShuffle + node.rng === nothing || + throw(TranslationError("`@shuffle` with an explicit rng cannot be pushed down to DuckDB, which has its own random number generator. Drop the `rng` argument, or materialize the query first.", :shuffle)) + push!(orderby_clauses, "random()") elseif node isa QueryableBackend.QueryableThenBy col = translate_orderby_expr(node.keySelector_expr, params) @@ -151,39 +162,99 @@ function build_sql(nodes::Vector{QueryableBackend.Queryable}, params::Vector{Any group_key_sql = key_col end - elseif node isa QueryableBackend.QueryableJoin - # Build JOIN clause - inner_source = node.inner - inner_from = if inner_source isa DuckDBQueryableSource - source_to_from(inner_source, "source_tbl_2") - else - throw(TranslationError("JOIN inner source must be a DuckDBQueryableSource", :join)) - end + elseif node isa QueryableBackend.QueryableJoin || + node isa QueryableBackend.QueryableLeftJoin || + node isa QueryableBackend.QueryableRightJoin || + node isa QueryableBackend.QueryableFullJoin + inner_from = inner_source_to_from(node, i, join_kind(node)) outer_alias = "t1" inner_alias = "t2" # Extract key selectors with table aliases - outer_key_sym, outer_key_body = extract_lambda_parts(node.outerKeySelector_expr) - outer_key_body = unwrap_block(outer_key_body) - outer_key_col = if is_property_access(outer_key_body, outer_key_sym) - quote_identifier(outer_alias) * "." * quote_identifier(extract_column_name(outer_key_body)) - else - translate_expr(outer_key_body, params, outer_key_sym) - end - inner_key_sym, inner_key_body = extract_lambda_parts(node.innerKeySelector_expr) - inner_key_body = unwrap_block(inner_key_body) - inner_key_col = if is_property_access(inner_key_body, inner_key_sym) - quote_identifier(inner_alias) * "." * quote_identifier(extract_column_name(inner_key_body)) - else - translate_expr(inner_key_body, params, inner_key_sym) - end - from_clause = "$from_clause AS $(quote_identifier(outer_alias)) INNER JOIN $inner_from AS $(quote_identifier(inner_alias)) ON $outer_key_col = $inner_key_col" + outer_key_col = join_key_sql(node.outerKeySelector_expr, outer_alias, params) + inner_key_col = join_key_sql(node.innerKeySelector_expr, inner_alias, params) + from_clause = "$from_clause AS $(quote_identifier(outer_alias)) $(join_kind(node)) $inner_from AS $(quote_identifier(inner_alias)) ON $outer_key_col = $inner_key_col" # Translate result selector with join context - if select_clause == "*" && hasproperty(node, :resultSelector_expr) + if select_clause == "*" select_clause = translate_join_map_expr(node.resultSelector_expr, params, outer_alias, inner_alias) end + elseif node isa QueryableBackend.QueryableConcat || + node isa QueryableBackend.QueryableUnion || + node isa QueryableBackend.QueryableExcept || + node isa QueryableBackend.QueryableIntersect + # Set operations combine the query built so far with a second one, + # so everything accumulated is sealed into the left-hand side and + # the result becomes the new FROM. + left_sql = assemble_sql(select_clause, from_clause, where_clauses, groupby_clauses, having_clauses, orderby_clauses, limit_clause, offset_clause, distinct, distinct_on) + right_sql = inner_query_sql(node, i, params) + + from_clause = setop_from(node, left_sql, right_sql, i) + + select_clause = "*" + where_clauses = String[] + having_clauses = String[] + orderby_clauses = String[] + groupby_clauses = String[] + limit_clause = nothing + offset_clause = nothing + distinct = false + distinct_on = nothing + has_groupby = false + group_key_sql = nothing + + elseif node isa QueryableBackend.QueryableCountBy + col = translate_groupby_expr(node.f_expr, params) + if select_clause != "*" || !isempty(groupby_clauses) + inner_sql = assemble_sql(select_clause, from_clause, where_clauses, groupby_clauses, having_clauses, orderby_clauses, limit_clause, offset_clause, distinct, distinct_on) + from_clause = "($inner_sql) AS countby_subq$(i)" + where_clauses = String[] + having_clauses = String[] + orderby_clauses = String[] + limit_clause = nothing + offset_clause = nothing + distinct = false + distinct_on = nothing + end + # The key column is named `key`, matching how the in-memory + # count_by and summarize name a scalar grouping key. + select_clause = "$col AS $(quote_identifier("key")), COUNT(*) AS $(quote_identifier("count"))" + groupby_clauses = [col] + has_groupby = false + group_key_sql = col + + elseif node isa QueryableBackend.QueryableTakeLast || + node isa QueryableBackend.QueryableDropLast + inner_sql = assemble_sql(select_clause, from_clause, where_clauses, groupby_clauses, having_clauses, orderby_clauses, limit_clause, offset_clause, distinct, distinct_on) + n = node.n + comparison = node isa QueryableBackend.QueryableTakeLast ? ">" : "<=" + # Everything accumulated is now inside inner_sql, so the clause + # state has to start over on top of it. + where_clauses = String[] + if n <= 0 + # take_last(0) keeps nothing; drop_last(0) keeps everything. + from_clause = "($inner_sql) AS lastsubq$(i)" + node isa QueryableBackend.QueryableTakeLast && push!(where_clauses, "FALSE") + else + # Row position is not a SQL concept, so it is materialised with + # ROW_NUMBER(). COUNT(*) OVER () supplies the total from the + # same single scan, so the subquery is not repeated and its + # positional parameters stay in order. + from_clause = "(SELECT * FROM ($inner_sql) AS lastinner$(i) " * + "QUALIFY ROW_NUMBER() OVER () $comparison COUNT(*) OVER () - $n) AS lastsubq$(i)" + end + select_clause = "*" + having_clauses = String[] + orderby_clauses = String[] + groupby_clauses = String[] + limit_clause = nothing + offset_clause = nothing + distinct = false + distinct_on = nothing + has_groupby = false + group_key_sql = nothing + else - throw(TranslationError("Unsupported query operation: $(typeof(node))", :unsupported)) + throw_unsupported(node) end end @@ -195,6 +266,76 @@ function build_sql(nodes::Vector{QueryableBackend.Queryable}, params::Vector{Any return SQLQuery(sql, params) end +# --- Two-input operators --- + +# Each two-input node registers its right-hand source under a name derived from +# its position in the walked tree, so that several of them in one query cannot +# collide. execution.jl walks the tree the same way to register them. +inner_table_name(i::Int) = "source_tbl_$(i)" + +join_kind(::QueryableBackend.QueryableJoin) = "INNER JOIN" +join_kind(::QueryableBackend.QueryableLeftJoin) = "LEFT OUTER JOIN" +join_kind(::QueryableBackend.QueryableRightJoin) = "RIGHT OUTER JOIN" +join_kind(::QueryableBackend.QueryableFullJoin) = "FULL OUTER JOIN" + +function inner_source_to_from(node, i::Int, label::AbstractString) + inner_source = node.inner + inner_source isa DuckDBQueryableSource || + throw(TranslationError("The second operand of $label must be a DuckDB source too — add `|> @duckdb()` to it.", :join)) + return source_to_from(inner_source, inner_table_name(i)) +end + +# A join key qualified by its table alias, so that a column present on both +# sides is unambiguous. +function join_key_sql(expr::Expr, alias::AbstractString, params::Vector{Any}) + sym, body = extract_lambda_parts(expr) + body = unwrap_block(body) + if is_property_access(body, sym) + return quote_identifier(alias) * "." * quote_identifier(extract_column_name(body)) + end + return translate_expr(body, params, sym) +end + +function inner_query_sql(node, i::Int, params::Vector{Any}) + inner_from = inner_source_to_from(node, i, "a set operation") + return "SELECT * FROM $inner_from" +end + +# The key of a `_by` set operation has to be a plain column: the key SQL is +# placed before the left-hand query in the generated text, so a key that +# contributed positional parameters would put them out of order. +function setop_key_sql(node, i::Int) + sym, body = extract_lambda_parts(node.f_expr) + body = unwrap_block(body) + is_property_access(body, sym) || + throw(TranslationError("A computed key in a `_by` set operation is not supported by the DuckDB backend; apply the transformation with `@map` first.", :setop)) + return quote_identifier(extract_column_name(body)) +end + +function setop_from(node::QueryableBackend.QueryableConcat, left_sql, right_sql, i::Int) + return "(($left_sql) UNION ALL ($right_sql)) AS setop$(i)" +end + +function setop_from(node::QueryableBackend.QueryableUnion, left_sql, right_sql, i::Int) + node.f_expr === nothing && return "(($left_sql) UNION ($right_sql)) AS setop$(i)" + key = setop_key_sql(node, i) + return "(SELECT DISTINCT ON ($key) * FROM (($left_sql) UNION ALL ($right_sql)) AS setopinner$(i)) AS setop$(i)" +end + +function setop_from(node::QueryableBackend.QueryableExcept, left_sql, right_sql, i::Int) + node.f_expr === nothing && return "(($left_sql) EXCEPT ($right_sql)) AS setop$(i)" + key = setop_key_sql(node, i) + return "(SELECT DISTINCT ON ($key) * FROM ($left_sql) AS setopleft$(i) " * + "WHERE $key NOT IN (SELECT $key FROM ($right_sql) AS setopright$(i))) AS setop$(i)" +end + +function setop_from(node::QueryableBackend.QueryableIntersect, left_sql, right_sql, i::Int) + node.f_expr === nothing && return "(($left_sql) INTERSECT ($right_sql)) AS setop$(i)" + key = setop_key_sql(node, i) + return "(SELECT DISTINCT ON ($key) * FROM ($left_sql) AS setopleft$(i) " * + "WHERE $key IN (SELECT $key FROM ($right_sql) AS setopright$(i))) AS setop$(i)" +end + function format_sql_option(value) if value isa AbstractString return "'$(escape_sql_string(value))'" diff --git a/src/unsupported.jl b/src/unsupported.jl new file mode 100644 index 0000000..907bec1 --- /dev/null +++ b/src/unsupported.jl @@ -0,0 +1,86 @@ +""" +Explanations for query operations the DuckDB backend cannot translate. + +`build_sql` used to fail with a bare "Unsupported query operation: ", +which told the user nothing about why or what to do instead. Each node that +has no SQL equivalent gets a method here saying what the obstacle is and how +to get the operation done anyway — almost always by materializing the query +first, which moves the rest of the pipeline onto the in-memory backend. +""" + +const _MATERIALIZE_HINT = "Materialize the query first, e.g. `… |> DataFrame |> " + +""" + unsupported_reason(node) -> String + +A human-readable explanation of why `node` cannot be pushed down to DuckDB. +""" +function unsupported_reason(node::QueryableBackend.Queryable) + return "the DuckDB backend cannot translate $(typeof(node)) into SQL." +end + +unsupported_reason(::QueryableBackend.QueryableChunk) = + "`@chunk` has no SQL equivalent — SQL has no batching construct. " * + _MATERIALIZE_HINT * "@chunk(3)`." + +unsupported_reason(::QueryableBackend.QueryableAggregateBy) = + "`@aggregate_by` folds each group with an arbitrary Julia function, which SQL " * + "cannot express. Use `@groupby` with `@map` and a SQL aggregate such as `sum` " * + "to push the aggregation down, or " * lowercasefirst(_MATERIALIZE_HINT) * "@aggregate_by(…)`." + +unsupported_reason(::QueryableBackend.QueryableTakeWhile) = + "`@take_while` stops at the first row failing the predicate, which depends on " * + "row order that SQL does not guarantee. Use `@filter` if you meant to keep every " * + "matching row, or " * lowercasefirst(_MATERIALIZE_HINT) * "@take_while(…)`." + +unsupported_reason(::QueryableBackend.QueryableDropWhile) = + "`@drop_while` skips a leading run of rows, which depends on row order that SQL " * + "does not guarantee. Use `@filter` if you meant to drop every matching row, or " * + lowercasefirst(_MATERIALIZE_HINT) * "@drop_while(…)`." + +unsupported_reason(::QueryableBackend.QueryableReverse) = + "`@reverse` has no SQL equivalent, because a SQL result has no inherent row order " * + "to reverse. Use `@orderby_descending` on the column you care about, or " * + lowercasefirst(_MATERIALIZE_HINT) * "@reverse()`." + +unsupported_reason(::QueryableBackend.QueryableIndex) = + "`@index` yields `(index, item)` pairs whose `item` is a whole row, and SQL has no " * + "nested row values. " * _MATERIALIZE_HINT * "@index()`." + +unsupported_reason(::QueryableBackend.QueryableAppend) = + "`@append` adds a single row at the end, which depends on row order that SQL does " * + "not guarantee. " * _MATERIALIZE_HINT * "@append(row)`." + +unsupported_reason(::QueryableBackend.QueryablePrepend) = + "`@prepend` adds a single row at the front, which depends on row order that SQL does " * + "not guarantee. " * _MATERIALIZE_HINT * "@prepend(row)`." + +unsupported_reason(::QueryableBackend.QueryableZip) = + "`@zip` pairs rows by position and yields tuples, which SQL has no value type for. " * + _MATERIALIZE_HINT * "@zip(other)`." + +unsupported_reason(::QueryableBackend.QueryableOfType) = + "`@of_type` dispatches on Julia types, which a SQL result does not carry. " * + _MATERIALIZE_HINT * "@of_type(T)`." + +unsupported_reason(::QueryableBackend.QueryableCast) = + "`@cast` converts to a Julia type, which SQL cannot express. Use `@map` with a " * + "conversion on individual columns, or " * lowercasefirst(_MATERIALIZE_HINT) * "@cast(T)`." + +unsupported_reason(::QueryableBackend.QueryableGroupJoin) = + "`@groupjoin` produces a nested collection per outer row, and SQL has no nested " * + "row values. Use `@join` with `@groupby`, or " * + lowercasefirst(_MATERIALIZE_HINT) * "@groupjoin(…)`." + +unsupported_reason(::QueryableBackend.QueryableMapMany) = + "`@mapmany` flattens a collection computed per row, which SQL cannot express in " * + "general. " * _MATERIALIZE_HINT * "@mapmany(…)`." + +""" + throw_unsupported(node) + +Raise a `TranslationError` carrying `unsupported_reason(node)`. +""" +function throw_unsupported(node::QueryableBackend.Queryable) + throw(TranslationError(unsupported_reason(node), :unsupported)) +end diff --git a/test/runtests.jl b/test/runtests.jl index 8c2459e..af0fcb2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,7 @@ using TestItemRunner include("test_queryduckdb.jl") +include("test_new_operators.jl") include("test_query_examples.jl") include("test_csvfiles.jl") include("test_excelfiles.jl") diff --git a/test/test_new_operators.jl b/test/test_new_operators.jl new file mode 100644 index 0000000..0c14253 --- /dev/null +++ b/test/test_new_operators.jl @@ -0,0 +1,303 @@ +@testsnippet NewOpData begin + using Query, QueryDuckDB, DataFrames, DataValues + import IteratorInterfaceExtensions + + # A Queryable is consumed through getiterator, not Base.collect, so this is + # how a non-table-shaped result is drained. + drain(q) = collect(IteratorInterfaceExtensions.getiterator(q)) + + people = DataFrame(id=[1, 2, 3], name=["John", "Sally", "Kirk"]) + pets = DataFrame(owner=[1, 3], pet=["Judy", "Ruff"]) + + nums = DataFrame(x=[1, 2, 2, 3]) + other = DataFrame(x=[3, 4]) + + # Set operations have no guaranteed row order in SQL, so results are + # compared as sorted vectors. + sorted(df, col) = sort(collect(skipmissing(df[!, col]))) +end + +@testitem "left join pushes down to a LEFT OUTER JOIN" setup=[NewOpData] begin + plan = people |> @duckdb() |> @left_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> @duckdbplan() + + @test occursin("LEFT OUTER JOIN", plan.sql) + + res = people |> @duckdb() |> @left_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> DataFrame + + @test size(res) == (3, 2) + @test sort(res[!, :name]) == ["John", "Kirk", "Sally"] + @test count(ismissing, res[!, :pet]) == 1 +end + +@testitem "right and full joins push down" setup=[NewOpData] begin + right_plan = people |> @duckdb() |> @right_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> @duckdbplan() + @test occursin("RIGHT OUTER JOIN", right_plan.sql) + + full_plan = people |> @duckdb() |> @full_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> @duckdbplan() + @test occursin("FULL OUTER JOIN", full_plan.sql) + + right = people |> @duckdb() |> @right_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> DataFrame + @test size(right) == (2, 2) + + full = people |> @duckdb() |> @full_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> DataFrame + @test size(full) == (3, 2) +end + +@testitem "outer joins agree with the in-memory backend" setup=[NewOpData] begin + duck = people |> @duckdb() |> @left_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> DataFrame + mem = people |> @left_join(pets, _.id, _.owner, {_.name, __.pet}) |> DataFrame + + @test sort(duck[!, :name]) == sort(mem[!, :name]) + @test sorted(duck, :pet) == sorted(mem, :pet) + @test count(ismissing, duck[!, :pet]) == count(ismissing, mem[!, :pet]) +end + +@testitem "an outer join needs a DuckDB source on both sides" setup=[NewOpData] begin + @test_throws QueryDuckDB.TranslationError begin + people |> @duckdb() |> @left_join(pets, _.id, _.owner, {_.name, __.pet}) |> DataFrame + end +end + +@testitem "concat, union, except and intersect push down" setup=[NewOpData] begin + @test occursin("UNION ALL", (nums |> @duckdb() |> @concat(other |> @duckdb()) |> @duckdbplan()).sql) + @test occursin("UNION", (nums |> @duckdb() |> @union(other |> @duckdb()) |> @duckdbplan()).sql) + @test occursin("EXCEPT", (nums |> @duckdb() |> @except(other |> @duckdb()) |> @duckdbplan()).sql) + @test occursin("INTERSECT", (nums |> @duckdb() |> @intersect(other |> @duckdb()) |> @duckdbplan()).sql) + + @test sorted(nums |> @duckdb() |> @concat(other |> @duckdb()) |> DataFrame, :x) == [1, 2, 2, 3, 3, 4] + @test sorted(nums |> @duckdb() |> @union(other |> @duckdb()) |> DataFrame, :x) == [1, 2, 3, 4] + @test sorted(nums |> @duckdb() |> @except(other |> @duckdb()) |> DataFrame, :x) == [1, 2] + @test sorted(nums |> @duckdb() |> @intersect(other |> @duckdb()) |> DataFrame, :x) == [3] +end + +@testitem "set operations agree with the in-memory backend" setup=[NewOpData] begin + @test sorted(nums |> @duckdb() |> @concat(other |> @duckdb()) |> DataFrame, :x) == + sorted(nums |> @concat(other) |> DataFrame, :x) + @test sorted(nums |> @duckdb() |> @union(other |> @duckdb()) |> DataFrame, :x) == + sorted(nums |> @union(other) |> DataFrame, :x) + @test sorted(nums |> @duckdb() |> @except(other |> @duckdb()) |> DataFrame, :x) == + sorted(nums |> @except(other) |> DataFrame, :x) + @test sorted(nums |> @duckdb() |> @intersect(other |> @duckdb()) |> DataFrame, :x) == + sorted(nums |> @intersect(other) |> DataFrame, :x) +end + +@testitem "the _by set operations push down to a key comparison" setup=[NewOpData] begin + a = DataFrame(k=[1, 2, 3], v=["a", "b", "c"]) + b = DataFrame(k=[2, 3], v=["B", "C"]) + + except_plan = a |> @duckdb() |> @except_by(b |> @duckdb(), _.k) |> @duckdbplan() + @test occursin("NOT IN", except_plan.sql) + + intersect_plan = a |> @duckdb() |> @intersect_by(b |> @duckdb(), _.k) |> @duckdbplan() + @test occursin("IN (SELECT", intersect_plan.sql) + + @test sorted(a |> @duckdb() |> @except_by(b |> @duckdb(), _.k) |> DataFrame, :k) == [1] + @test sorted(a |> @duckdb() |> @intersect_by(b |> @duckdb(), _.k) |> DataFrame, :k) == [2, 3] + @test sorted(a |> @duckdb() |> @union_by(b |> @duckdb(), _.k) |> DataFrame, :k) == [1, 2, 3] +end + +@testitem "a computed key in a _by set operation is rejected with an explanation" setup=[NewOpData] begin + a = DataFrame(k=[1, 2], v=["a", "b"]) + b = DataFrame(k=[2], v=["B"]) + + err = try + a |> @duckdb() |> @except_by(b |> @duckdb(), _.k * 2) |> DataFrame + nothing + catch e + e + end + + @test err isa QueryDuckDB.TranslationError + @test occursin("@map", err.msg) +end + +@testitem "order and order_descending push down to ORDER BY ALL" setup=[NewOpData] begin + asc = nums |> @duckdb() |> @order() |> @duckdbplan() + @test occursin("ORDER BY ALL ASC", asc.sql) + + desc = nums |> @duckdb() |> @order_descending() |> @duckdbplan() + @test occursin("ORDER BY ALL DESC", desc.sql) + + @test (nums |> @duckdb() |> @order() |> DataFrame)[!, :x] == [1, 2, 2, 3] + @test (nums |> @duckdb() |> @order_descending() |> DataFrame)[!, :x] == [3, 2, 2, 1] +end + +@testitem "orderby with a real key selector is unaffected by order" setup=[NewOpData] begin + plan = people |> @duckdb() |> @orderby(_.name) |> @duckdbplan() + + @test occursin("ORDER BY", plan.sql) + @test !occursin("ORDER BY ALL", plan.sql) + @test (people |> @duckdb() |> @orderby(_.name) |> DataFrame)[!, :name] == ["John", "Kirk", "Sally"] +end + +@testitem "shuffle pushes down to ORDER BY random()" setup=[NewOpData] begin + plan = nums |> @duckdb() |> @shuffle() |> @duckdbplan() + @test occursin("random()", plan.sql) + + @test sort((nums |> @duckdb() |> @shuffle() |> DataFrame)[!, :x]) == [1, 2, 2, 3] +end + +@testitem "shuffle with an explicit rng cannot be pushed down" setup=[NewOpData] begin + using Random + + err = try + nums |> @duckdb() |> @shuffle(rng=MersenneTwister(1)) |> DataFrame + nothing + catch e + e + end + + @test err isa QueryDuckDB.TranslationError + @test occursin("rng", err.msg) +end + +@testitem "take_last and drop_last push down with a window function" setup=[NewOpData] begin + df = DataFrame(x=[1, 2, 3, 4, 5]) + + plan = df |> @duckdb() |> @take_last(2) |> @duckdbplan() + @test occursin("ROW_NUMBER() OVER ()", plan.sql) + @test occursin("QUALIFY", plan.sql) + + @test (df |> @duckdb() |> @take_last(2) |> DataFrame)[!, :x] == [4, 5] + @test (df |> @duckdb() |> @drop_last(2) |> DataFrame)[!, :x] == [1, 2, 3] + + # A count of zero or less behaves as it does in memory. + @test nrow(df |> @duckdb() |> @take_last(0) |> DataFrame) == 0 + @test (df |> @duckdb() |> @drop_last(0) |> DataFrame)[!, :x] == [1, 2, 3, 4, 5] +end + +@testitem "take_last and drop_last agree with the in-memory backend" setup=[NewOpData] begin + df = DataFrame(x=[1, 2, 3, 4, 5]) + + @test (df |> @duckdb() |> @take_last(2) |> DataFrame) == (df |> @take_last(2) |> DataFrame) + @test (df |> @duckdb() |> @drop_last(2) |> DataFrame) == (df |> @drop_last(2) |> DataFrame) +end + +@testitem "count_by pushes down to GROUP BY with COUNT" setup=[NewOpData] begin + df = DataFrame(k=["a", "b", "a"], v=[1, 2, 3]) + + plan = df |> @duckdb() |> @count_by(_.k) |> @duckdbplan() + @test occursin("COUNT(*)", plan.sql) + @test occursin("GROUP BY", plan.sql) + + duck = df |> @duckdb() |> @count_by(_.k) |> DataFrame + mem = df |> @count_by(_.k) |> DataFrame + + @test names(duck) == ["key", "count"] + @test sort(duck, :key) == sort(mem, :key) +end + +@testitem "terminal operators push down into the SQL" setup=[NewOpData] begin + df = DataFrame(x=[1, 2, 3, 4]) + + @test (df |> @duckdb() |> @count()) == 4 + @test (df |> @duckdb() |> @filter(_.x > 2) |> @count()) == 2 + @test (df |> @duckdb() |> @any()) == true + @test @any(df |> @duckdb(), _.x > 3) == true + @test @any(df |> @duckdb(), _.x > 9) == false + @test (df |> @duckdb() |> @all(_.x > 0)) == true + @test (df |> @duckdb() |> @all(_.x > 1)) == false + + @test (df |> @duckdb() |> @first()).x == 1 + @test (df |> @duckdb() |> @element_at(2)).x == 2 + @test (df |> @duckdb() |> @min_by(_.x)).x == 1 + @test (df |> @duckdb() |> @max_by(_.x)).x == 4 +end + +@testitem "terminal operators agree with the in-memory backend" setup=[NewOpData] begin + df = DataFrame(x=[3, 1, 4, 1, 5]) + + @test (df |> @duckdb() |> @count()) == (df |> @count()) + @test (df |> @duckdb() |> @any()) == (df |> @any()) + @test (df |> @duckdb() |> @all(_.x > 0)) == (df |> @all(_.x > 0)) + @test (df |> @duckdb() |> @min_by(_.x)).x == (df |> @min_by(_.x)).x + @test (df |> @duckdb() |> @max_by(_.x)).x == (df |> @max_by(_.x)).x +end + +@testitem "count on a DuckDB query used to have no method" setup=[NewOpData] begin + # Regression: QueryOperators.count had no Queryable method before + # QueryableScalar, so this call failed outright. + @test (people |> @duckdb() |> @filter(_.id > 1) |> @count()) == 2 +end + +@testitem "terminal operators with no SQL translation fall back to memory" setup=[NewOpData] begin + df = DataFrame(x=[1, 2, 3]) + + # aggregate, last, single, contains and sequence_equal are not translated, + # so they materialize and run the in-memory implementation. + @test (df |> @duckdb() |> @aggregate((acc, cur) -> (x = acc.x + cur.x,))).x == 6 + @test (df |> @duckdb() |> @last()).x == 3 + @test @single(df |> @duckdb(), _.x == 2).x == 2 + @test (df |> @duckdb() |> @contains((x=2,))) == true +end + +@testitem "terminal operators report empty sequences the same way" setup=[NewOpData] begin + empty = DataFrame(x=Int[]) + + @test (empty |> @duckdb() |> @count()) == 0 + @test (empty |> @duckdb() |> @any()) == false + @test_throws ErrorException (empty |> @duckdb() |> @first()) + @test_throws ErrorException (empty |> @duckdb() |> @min_by(_.x)) + @test_throws ErrorException (empty |> @duckdb() |> @element_at(1)) +end + +@testitem "operators with no SQL equivalent explain themselves" setup=[NewOpData] begin + df = DataFrame(x=[1, 2, 3]) + + cases = [ + (() -> drain(df |> @duckdb() |> @chunk(2)), "@chunk"), + (() -> df |> @duckdb() |> @reverse() |> DataFrame, "@reverse"), + (() -> drain(df |> @duckdb() |> @index()), "@index"), + (() -> df |> @duckdb() |> @take_while(_.x < 3) |> DataFrame, "@take_while"), + (() -> df |> @duckdb() |> @drop_while(_.x < 3) |> DataFrame, "@drop_while"), + (() -> df |> @duckdb() |> @append((x=4,)) |> DataFrame, "@append"), + (() -> df |> @duckdb() |> @prepend((x=0,)) |> DataFrame, "@prepend"), + (() -> drain(df |> @duckdb() |> @of_type(NamedTuple)), "@of_type"), + (() -> drain(df |> @duckdb() |> @cast(Any)), "@cast"), + (() -> df |> @duckdb() |> @aggregate_by(_.x, 0, (a, c) -> a + c.x) |> DataFrame, "@aggregate_by"), + ] + + for (f, name) in cases + err = try + f() + nothing + catch e + e + end + + @test err isa QueryDuckDB.TranslationError + # The message names the operator and suggests a way forward, rather + # than just printing the node type. + @test occursin(name, err.msg) + @test occursin("Materialize", err.msg) || occursin("materialize", err.msg) || occursin("Use ", err.msg) + end +end + +@testitem "DuckDB rows carry DataValue, not missing" setup=[NewOpData] begin + rows = drain(people |> @duckdb() |> @left_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet})) + + @test length(rows) == 3 + @test all(r -> r.pet isa DataValue, rows) + @test !any(r -> ismissing(r.pet), rows) + @test count(r -> isna(r.pet), rows) == 1 +end + +@testitem "a DuckDB result still lands in a DataFrame as missing" setup=[NewOpData] begin + # The column interface still speaks Missing, which is what TableTraits' + # _using_missing protocol and table sinks expect. + res = people |> @duckdb() |> @left_join(pets |> @duckdb(), _.id, _.owner, {_.name, __.pet}) |> DataFrame + + @test count(ismissing, res[!, :pet]) == 1 + @test eltype(res[!, :pet]) == Union{Missing,String} +end + +@testitem "a DuckDB result can be piped back through query operators" setup=[NewOpData] begin + res = people |> @duckdb() |> @filter(_.id > 1) |> @map({_.name}) + + # Consuming a DuckDB result through the in-memory operators agrees with + # running the whole thing in memory. + again = res |> @order() |> DataFrame + mem = people |> @filter(_.id > 1) |> @map({_.name}) |> @order() |> DataFrame + + @test again == mem +end