Push the new LINQ operators down into SQL

Project.toml

@@ -1,10 +1,11 @@
 name = "QueryDuckDB"
 uuid = "fb67597e-2b1e-4ffa-9078-014b5866405e"
 version = "1.0.0-DEV"
 authors = ["David Anthoff <anthoff@berkeley.edu>"]
 
 [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 @@
 [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"

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.

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

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

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 ---

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