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