beam-duckdb-0.3.1.0: tests/Database/Beam/DuckDB/Test/Query.hs
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-type-defaults #-}
module Database.Beam.DuckDB.Test.Query (tests) where
import Control.Monad (void)
import Data.Int (Int32)
import Data.List (nub, nubBy, sort, sortOn)
import Data.Text (Text)
import Database.Beam
( Beamable,
Columnar,
Database,
DatabaseSettings,
Generic,
Identity,
SqlEq ((==.)),
SqlOrd ((>.)),
SqlValable (val_),
Table (..),
TableEntity,
aggregate_,
allInGroup_,
all_,
anyOver_,
as_,
asc_,
avgOver_,
bounds_,
cast_,
concat_,
countAll_,
dbModification,
defaultDbSettings,
desc_,
double,
everyOver_,
filterWhere_,
filter_,
firstValue_,
frame_,
fromMaybe_,
group_,
guard_,
insert,
insertValues,
isTrue_,
lagWithDefault_,
lastValue_,
lead1_,
leftJoin_,
modifyTableFields,
noBounds_,
noOrder_,
noPartition_,
nrows_,
nub_,
orderBy_,
orderPartitionBy_,
over_,
partitionBy_,
references_,
related_,
reuse,
rowNumber_,
runInsert,
runSelectReturningList,
select,
selectWith,
selecting,
someOver_,
sqlBool_,
sumOver_,
sum_,
tableModification,
unbounded_,
union_,
varchar,
withDbModification,
withWindow_,
(<.),
(>=.),
)
import Database.Beam.DuckDB (DuckDB, runBeamDuckDB)
import Database.DuckDB.Simple (Connection, execute_, withConnection)
import Hedgehog (Gen, annotate, evalIO, forAll, property, (===))
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (testCase, (@?=))
import Test.Tasty.Hedgehog (testProperty)
tests :: TestTree
tests =
testGroup
"Queries"
[ testGroup
"SQL92 featureset"
[ testGroup
"Selection"
[ testSelectAll,
testSelectWithFilter,
testSelectEquality
],
testGroup
"Projection"
[ testProjectSingleColumn,
testProjectMultipleColumns,
testProjectWithExpression
],
testGroup
"Join"
[ testInnerJoin,
testMultiInnerJoin,
testMixingJoinsWithFilters,
testLeftJoin
]
],
testGroup
"SQL99 featureset"
[ testGroup "CONCAT" [testConcat],
testGroup
"CTE"
[ testCommonTableExpression,
testMultipleCommonTableExpressions,
testRecursiveCommonTableExpression
],
testGroup
"SOME/ANY/EVERY"
[ testSomeOver,
testAnyOver,
testEveryOver
]
],
testGroup
"SQL2003 featureset"
[ testGroup
"Windowing"
[ testRowNumberOverWholeResult,
testRowNumberPartitionedByColumn,
testWindowingWithBounds
],
testGroup
"LAG/LEAD"
[ testLag,
testLead,
testLeadOverNub
],
testGroup
"FILTER"
[ testFilterWhereCountAll,
testFilterWithWindow
],
testGroup
"FIRST_VALUE/LAST_VALUE"
[ testFirstValue,
testLastValue
]
]
]
testSelectAll :: TestTree
testSelectAll = testProperty "selecting all users should return the users initially inserted" $ property $ do
users <- forAll genUsers
results <- evalIO $
withTestDb users [] [] $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select (all_ (_dbUsers testDb))
sortOn _userId results === sortOn _userId users
testSelectWithFilter :: TestTree
testSelectWithFilter = testProperty "selecting users satisfying a condition works as expected" $ property $ do
users <- forAll genUsers
ageThreshold <- forAll genAge
results <- evalIO $
withTestDb users [] [] $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
u <- all_ (_dbUsers testDb)
guard_ (_userAge u >. val_ ageThreshold)
pure u
let expected = filter (\u -> _userAge u > ageThreshold) users
sortOn _userId results === sortOn _userId expected
testSelectEquality :: TestTree
testSelectEquality = testProperty "" $ property $ do
users <- forAll genUsers
target <- forAll (Gen.element users)
results <- evalIO $
withTestDb users [] [] $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
u <- all_ (_dbUsers testDb)
guard_ (_userId u ==. val_ (_userId target))
pure u
results === [target]
testProjectSingleColumn :: TestTree
testProjectSingleColumn = testProperty "Single column projection works as expected" $ property $ do
users <- forAll genUsers
results <- evalIO $
withTestDb users [] [] $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
u <- all_ (_dbUsers testDb)
pure (_userName u)
sort results === sort (map _userName users)
testProjectMultipleColumns :: TestTree
testProjectMultipleColumns = testProperty "Multiple column projection works as expected" $ property $ do
users <- forAll genUsers
results <- evalIO $
withTestDb users [] [] $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
u <- all_ (_dbUsers testDb)
pure (_userName u, _userAge u)
let expected = map (\u -> (_userName u, _userAge u)) users
sort results === sort expected
testProjectWithExpression :: TestTree
testProjectWithExpression = testProperty "Projection using an expression works as expected" $ property $ do
products <- forAll genProducts
results <- evalIO $
withTestDb [] products [] $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
p <- all_ (_dbProducts testDb)
pure (_productId p, _productPrice p * val_ 2)
let expected = map (\p -> (_productId p, _productPrice p * 2)) products
sortOn fst results === sortOn fst expected
testInnerJoin :: TestTree
testInnerJoin = testProperty "Inner joins work as expected" $ property $ do
users <- forAll genUsers
products <- forAll genProducts
orders <- forAll (genOrders users products)
results <- evalIO $
withTestDb users products orders $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
o <- all_ (_dbOrders testDb)
u <- related_ (_dbUsers testDb) (_orderUserId o)
pure (_userName u, _orderQuantity o)
let expected = do
o <- orders
let UserId uid = _orderUserId o
u <- filter (\u -> _userId u == uid) users
pure (_userName u, _orderQuantity o)
sort results === sort expected
testMultiInnerJoin :: TestTree
testMultiInnerJoin = testProperty "Multi-inner joins work as expected" $ property $ do
users <- forAll genUsers
products <- forAll genProducts
orders <- forAll (genOrders users products)
results <- evalIO $
withTestDb users products orders $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
o <- all_ (_dbOrders testDb)
u <- related_ (_dbUsers testDb) (_orderUserId o)
p <- related_ (_dbProducts testDb) (_orderProductId o)
pure (_userName u, _productName p, _orderQuantity o)
let expected = do
o <- orders
let UserId uid = _orderUserId o
ProductId pid = _orderProductId o
u <- filter (\u -> _userId u == uid) users
p <- filter (\p -> _productId p == pid) products
pure (_userName u, _productName p, _orderQuantity o)
sort results === sort expected
testMixingJoinsWithFilters :: TestTree
testMixingJoinsWithFilters = testProperty "Mixing joins and filters works as expected" $ property $ do
users <- forAll genUsers
products <- forAll genProducts
orders <- forAll (genOrders users products)
minQty <- forAll (Gen.int32 (Range.linear 1 50))
results <- evalIO $
withTestDb users products orders $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
o <- all_ (_dbOrders testDb)
guard_ (_orderQuantity o >=. val_ minQty)
u <- related_ (_dbUsers testDb) (_orderUserId o)
pure (_userName u, _orderQuantity o)
let expected = do
o <- filter (\o -> _orderQuantity o >= minQty) orders
let UserId uid = _orderUserId o
u <- filter (\u -> _userId u == uid) users
pure (_userName u, _orderQuantity o)
sort results === sort expected
testLeftJoin :: TestTree
testLeftJoin = testProperty "Left joins work as expected" $ property $ do
users <- forAll genUsers
products <- forAll genProducts
-- Generate orders for only a subset of users
let halfUsers = take (length users `div` 2) users
orders <- forAll (genOrders halfUsers products)
results <- evalIO $
withTestDb users products orders $ \conn ->
runBeamDuckDB conn $
runSelectReturningList $
select $ do
u <- all_ (_dbUsers testDb)
o <-
leftJoin_
(all_ (_dbOrders testDb))
(\o -> _orderUserId o ==. primaryKey u)
pure (_userId u, _orderQuantity o)
let resultUserIds = nubBy (\a b -> fst a == fst b) (sortOn fst results)
length resultUserIds === length users
let usersWithoutOrders = filter (\u -> _userId u `notElem` map _userId halfUsers) users
nothingRows = filter (\(uid, _) -> uid `elem` map _userId usersWithoutOrders) results
annotate "Users without orders should have Nothing quantity"
mapM_ (\(_, mq) -> mq === Nothing) nothingRows
-- TODO: this generates an invalid query
-- testCountDistinct :: TestTree
-- testCountDistinct = testCase "counts distinct users who placed orders" $ do
-- let users = [User 1 "Alice" 30]
-- withTestDb users [] [] $ \conn -> do
-- result <-
-- runBeamDuckDB conn $
-- runSelectReturningList $
-- select $ do
-- order <- all_ (_dbOrders testDb)
-- guard_ (distinct_ (pure $ _orderUserId order))
-- pure order
-- result @?= []
testConcat :: TestTree
testConcat = testCase "CONCAT two columns" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
withTestDb users [] [] $ \conn -> do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
orderBy_ (asc_ . fst) $
do
user <- all_ (_dbUsers testDb)
pure
( _userId user,
concat_
[ _userName user,
val_ " (age: ",
cast_ (_userAge user) (varchar Nothing),
val_ ")"
]
)
map snd rows
@?= [ "Alice (age: 30)",
"Bob (age: 25)",
"Charlie (age: 35)"
]
testCommonTableExpression :: TestTree
testCommonTableExpression = testCase "Non-recursive common table expression" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
selectWith $ do
userTotals <-
selecting
$ aggregate_
( \o ->
( group_ (_orderUserId o),
fromMaybe_ (val_ 0) $ sum_ (_orderQuantity o)
)
)
$ all_ (_dbOrders testDb)
pure $ do
(uid, total) <- reuse userTotals
u <- filter_ (\u -> UserId (_userId u) ==. uid) $ all_ (_dbUsers testDb)
pure (_userName u, total)
rows @?= [("Alice", 3), ("Bob", 1)]
testMultipleCommonTableExpressions :: TestTree
testMultipleCommonTableExpressions = testCase "Multiple common table expressions" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn -> do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
selectWith $ do
-- CTE 1: expensive products (price > 1000 cents)
expensiveProducts <-
selecting $
filter_ (\p -> _productPrice p >. val_ 1000) $
all_ (_dbProducts testDb)
-- CTE 2: orders for expensive products
expensiveOrders <-
selecting $ do
p <- reuse expensiveProducts
filter_ (\o -> _orderProductId o `references_` p) $
all_ (_dbOrders testDb)
-- Main query: distinct users who ordered expensive products
pure $ do
o <- reuse expensiveOrders
u <-
filter_ (\u -> _orderUserId o `references_` u) $
all_ (_dbUsers testDb)
pure (_userName u)
-- Gadget costs 2999; Alice ordered it
nub rows @?= ["Alice"]
testRecursiveCommonTableExpression :: TestTree
testRecursiveCommonTableExpression = testCase "Recursive common table expression" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
selectWith $ do
rec cte <-
selecting
-- Base case: SELECT 1
( pure (as_ @Int32 (val_ 1))
`union_`
-- Recursive step: SELECT n + 1 FROM cte WHERE n < 5
( do
n <- reuse cte
guard_ (n <. val_ 5)
pure (n + val_ 1)
)
)
pure (reuse cte)
sort rows @?= [1, 2, 3, 4, 5 :: Int32]
testSomeOver :: TestTree
testSomeOver = testCase "SOME" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
aggregate_
( \o ->
-- Expecting that at least one order per user has a quantity larger than 1
-- This is only true of a single user, UserId 1
( group_ (_orderUserId o),
isTrue_ $ someOver_ allInGroup_ (sqlBool_ (_orderQuantity o >. val_ 1))
)
)
(all_ (_dbOrders testDb))
rows
@?= [ (UserId 1, True),
(UserId 2, False)
]
testAnyOver :: TestTree
testAnyOver = testCase "ANY" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
aggregate_
( \o ->
-- Expecting that at least one order per user has a quantity larger than 1
-- This is only true of a single user, UserId 1
( group_ (_orderUserId o),
isTrue_ $ anyOver_ allInGroup_ (sqlBool_ (_orderQuantity o >. val_ 1))
)
)
(all_ (_dbOrders testDb))
rows
@?= [ (UserId 1, True),
(UserId 2, False)
]
testEveryOver :: TestTree
testEveryOver = testCase "EVERY" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
aggregate_
( \o ->
-- Expecting that all orders per user have a quantity larger than 1
-- This is not true for any user
( group_ (_orderUserId o),
isTrue_ $ everyOver_ allInGroup_ (sqlBool_ (_orderQuantity o >. val_ 1))
)
)
(all_ (_dbOrders testDb))
rows
@?= [ (UserId 1, False),
(UserId 2, False)
]
testRowNumberOverWholeResult :: TestTree
testRowNumberOverWholeResult = testCase "ROW_NUMBERS() over entire result set" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
withTestDb users [] [] $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
withWindow_
(\_ -> frame_ noPartition_ noOrder_ noBounds_)
(\u w -> (_userName u, rowNumber_ `over_` w))
(all_ (_dbUsers testDb))
rows
@?= [ ("Alice", 1 :: Int32),
("Bob", 2),
("Charlie", 3)
]
testRowNumberPartitionedByColumn :: TestTree
testRowNumberPartitionedByColumn = testCase "ROW_NUMBERS() partitioned over column" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 3
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
orderBy_ (\(oid, _, _, _) -> asc_ oid) $
withWindow_
( \o ->
frame_
(partitionBy_ (_orderProductId o))
(orderPartitionBy_ (desc_ (_orderQuantity o)))
noBounds_
)
(\o w -> (_orderId o, _orderProductId o, _orderUserId o, as_ @Int32 (rowNumber_ `over_` w)))
(all_ (_dbOrders testDb))
-- We expect two groups of rows.
-- One group of rows for product ID 1, one group of rows for product ID 2.
-- For each group, the data for orders is considered in descending number of quantity
--
-- the order of each group is not deterministic.
rows @?= [(1, ProductId 1, UserId 1, 1), (2, ProductId 1, UserId 2, 2), (3, ProductId 2, UserId 1, 1)]
testWindowingWithBounds :: TestTree
testWindowingWithBounds = testCase "3-row moving average" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25
]
products =
[Product 1 "Widget" 999]
orders =
[ Order 1 (UserId 1) (ProductId 1) 10,
Order 2 (UserId 1) (ProductId 1) 5,
Order 3 (UserId 2) (ProductId 1) 8,
Order 4 (UserId 1) (ProductId 1) 3,
Order 5 (UserId 2) (ProductId 1) 7
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
withWindow_
( \o ->
frame_
noPartition_
(orderPartitionBy_ (asc_ (_orderId o)))
(bounds_ (nrows_ 1) (Just (nrows_ 1)))
)
( \o w ->
( _orderId o,
avgOver_ allInGroup_ (cast_ (_orderQuantity o) double) `over_` w
)
)
(all_ (_dbOrders testDb))
rows
@?= [ (1, Just 7.5),
(2, Just 7.666666666666667),
(3, Just 5.333333333333333),
(4, Just 6.0),
(5, Just 5.0)
]
testLag :: TestTree
testLag = testCase "LAG" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 3
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
withWindow_
(\o -> frame_ noPartition_ (orderPartitionBy_ (asc_ (_orderId o))) noBounds_)
( \o w ->
lagWithDefault_ (_orderQuantity o) (val_ (1 :: Int32)) (val_ 0) `over_` w
)
(all_ (_dbOrders testDb))
-- First row has no predecessor → default 0
-- Second row lags to first row's qty (2)
-- Third row lags to second row's qty (1)
rows @?= [0, 2, 1]
testLead :: TestTree
testLead = testCase "LEAD" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 3
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
withWindow_
(\o -> frame_ noPartition_ (orderPartitionBy_ (asc_ (_orderId o))) noBounds_)
( \o w ->
lagWithDefault_ (_orderQuantity o) (val_ (1 :: Int32)) (val_ 0) `over_` w
)
(all_ (_dbOrders testDb))
-- First leads to second's qty (1)
-- Second leads to third's qty (1)
-- Third has no successor → default 0
rows @?= [0, 2, 1]
-- Regression test for #746
testLeadOverNub :: TestTree
testLeadOverNub = testCase "LEAD over nub_ (issue #746)" $ do
let users = [User 1 "Alice" 30]
products = [Product 1 "Widget" 999]
orders =
[ Order 1 (UserId 1) (ProductId 1) 10,
Order 2 (UserId 1) (ProductId 1) 10,
Order 3 (UserId 1) (ProductId 1) 20,
Order 4 (UserId 1) (ProductId 1) 20,
Order 5 (UserId 1) (ProductId 1) 30,
Order 6 (UserId 1) (ProductId 1) 30
]
withTestDb users products orders $ \conn -> do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
withWindow_
(\q -> frame_ noPartition_ (orderPartitionBy_ (asc_ q)) noBounds_)
(\q w -> (q, lead1_ q `over_` w))
(nub_ (_orderQuantity <$> all_ (_dbOrders testDb)))
rows
@?= [ (10 :: Int32, Just 20),
(20, Just 30),
(30, Nothing)
]
testFilterWhereCountAll :: TestTree
testFilterWhereCountAll = testCase "COUNT(*) FILTER (WHERE ... )" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 3
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn
$ runSelectReturningList
$ select
$ aggregate_
( \o ->
( group_ (_orderUserId o),
as_ @Int32 $ countAll_,
as_ @Int32 $ countAll_ `filterWhere_` (_orderQuantity o >. val_ 1)
)
)
$ all_ (_dbOrders testDb)
-- Alice (uid 1): 2 total orders, 2 with qty > 1
-- Bob (uid 2): 1 total order, 0 with qty > 1
rows @?= [(UserId 1, 2, 2), (UserId 2, 1, 0)]
testFilterWithWindow :: TestTree
testFilterWithWindow = testCase "FILTER with window function" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
withWindow_
(\o -> frame_ noPartition_ (orderPartitionBy_ (asc_ (_orderId o))) noBounds_)
( \o w ->
( _orderId o,
fromMaybe_
(val_ 0)
( ( sumOver_ allInGroup_ (_orderQuantity o)
`filterWhere_` (_orderQuantity o >. val_ 1)
)
`over_` w
)
)
)
(all_ (_dbOrders testDb))
-- Only order 1 (qty=2) passes the filter
-- Running filtered sum: 2, 2, 2
let sorted = sort rows
map snd sorted @?= [2, 2, 2]
testFirstValue :: TestTree
testFirstValue = testCase "FIRST_VALUE" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
orderBy_ (\(oid, _, _) -> desc_ oid) $
withWindow_
( \o ->
frame_
(partitionBy_ (_orderProductId o))
(orderPartitionBy_ (asc_ (_orderQuantity o)))
(bounds_ unbounded_ (Just unbounded_))
)
( \o w ->
( _orderId o,
_orderProductId o,
firstValue_ (_orderQuantity o) `over_` w
)
)
(all_ (_dbOrders testDb))
-- Widget (pid 1): quantities 1, 2 → first_value is 1 for both
-- Gadget (pid 2): quantity 1 → first_value is 1
rows @?= [(3, ProductId 2, 1), (2, ProductId 1, 1), (1, ProductId 1, 1)]
testLastValue :: TestTree
testLastValue = testCase "LAST_VALUE" $ do
let users =
[ User 1 "Alice" 30,
User 2 "Bob" 25,
User 3 "Charlie" 35
]
products =
[ Product 1 "Widget" 999,
Product 2 "Gadget" 2999
]
orders =
[ Order 1 (UserId 1) (ProductId 1) 2,
Order 2 (UserId 2) (ProductId 1) 1,
Order 3 (UserId 1) (ProductId 2) 1
]
withTestDb users products orders $ \conn ->
do
rows <-
runBeamDuckDB conn $
runSelectReturningList $
select $
orderBy_ (\(oid, _, _) -> asc_ oid) $
withWindow_
( \o ->
frame_
(partitionBy_ (_orderProductId o))
(orderPartitionBy_ (asc_ (_orderQuantity o)))
-- Must use UNBOUNDED FOLLOWING to see the actual last value,
-- otherwise the default frame ends at CURRENT ROW.
(bounds_ unbounded_ (Just unbounded_))
)
( \o w ->
( _orderId o,
_orderProductId o,
lastValue_ (_orderQuantity o) `over_` w
)
)
(all_ (_dbOrders testDb))
-- Widget (pid 1): quantities 1, 2 → last_value is 2 for both
rows @?= [(1, ProductId 1, 2), (2, ProductId 1, 2), (3, ProductId 2, 1)]
data UserT f = User
{ _userId :: Columnar f Int32,
_userName :: Columnar f Text,
_userAge :: Columnar f Int32
}
deriving (Generic)
type User = UserT Identity
type UserId = PrimaryKey UserT Identity
deriving instance Show UserId
deriving instance Eq UserId
deriving instance Ord UserId
deriving instance Show User
deriving instance Eq User
deriving instance Ord User
instance Beamable UserT
instance Table UserT where
data PrimaryKey UserT f = UserId (Columnar f Int32)
deriving (Generic)
primaryKey = UserId . _userId
instance Beamable (PrimaryKey UserT)
data ProductT f = Product
{ _productId :: Columnar f Int32,
_productName :: Columnar f Text,
_productPrice :: Columnar f Int32 -- cents, to avoid floating point
}
deriving (Generic)
type Product = ProductT Identity
type ProductId = PrimaryKey ProductT Identity
deriving instance Show ProductId
deriving instance Eq ProductId
deriving instance Ord ProductId
deriving instance Show Product
deriving instance Eq Product
deriving instance Ord Product
instance Beamable ProductT
instance Table ProductT where
data PrimaryKey ProductT f = ProductId (Columnar f Int32)
deriving (Generic)
primaryKey = ProductId . _productId
instance Beamable (PrimaryKey ProductT)
data OrderT f = Order
{ _orderId :: Columnar f Int32,
_orderUserId :: PrimaryKey UserT f,
_orderProductId :: PrimaryKey ProductT f,
_orderQuantity :: Columnar f Int32
}
deriving (Generic)
type Order = OrderT Identity
type OrderId = PrimaryKey OrderT Identity
deriving instance Show OrderId
deriving instance Eq OrderId
deriving instance Ord OrderId
deriving instance Show Order
deriving instance Eq Order
deriving instance Ord Order
instance Beamable OrderT
instance Table OrderT where
data PrimaryKey OrderT f = OrderId (Columnar f Int32)
deriving (Generic)
primaryKey = OrderId . _orderId
instance Beamable (PrimaryKey OrderT)
data TestDB f = TestDB
{ _dbUsers :: f (TableEntity UserT),
_dbProducts :: f (TableEntity ProductT),
_dbOrders :: f (TableEntity OrderT)
}
deriving (Generic, Database be)
testDb :: DatabaseSettings DuckDB TestDB
testDb =
defaultDbSettings
`withDbModification` dbModification
{ _dbUsers =
modifyTableFields
tableModification
{ _userId = "id",
_userName = "name",
_userAge = "age"
},
_dbProducts =
modifyTableFields
tableModification
{ _productId = "id",
_productName = "name",
_productPrice = "price"
},
_dbOrders =
modifyTableFields
tableModification
{ _orderId = "id",
_orderUserId = UserId "user_id",
_orderProductId = ProductId "product_id",
_orderQuantity = "quantity"
}
}
genName :: Gen Text
genName = Gen.text (Range.linear 1 50) Gen.alphaNum
genAge :: Gen Int32
genAge = Gen.int32 (Range.linear 1 120)
genPrice :: Gen Int32
genPrice = Gen.int32 (Range.linear 100 100000)
genQuantity :: Gen Int32
genQuantity = Gen.int32 (Range.linear 1 100)
genUsers :: Gen [User]
genUsers = do
n <- Gen.int (Range.linear 3 20)
traverse (\i -> User (fromIntegral i) <$> genName <*> genAge) [1 .. n]
genProducts :: Gen [Product]
genProducts = do
n <- Gen.int (Range.linear 2 10)
traverse (\i -> Product (fromIntegral i) <$> genName <*> genPrice) [1 .. n]
genOrders :: [User] -> [Product] -> Gen [Order]
genOrders users products = do
n <- Gen.int (Range.linear 1 (length users * length products))
traverse
( \i -> do
uid <- Gen.element (map _userId users)
pid <- Gen.element (map _productId products)
Order (fromIntegral i) (UserId uid) (ProductId pid) <$> genQuantity
)
[1 .. n]
createTables :: Connection -> IO ()
createTables conn = do
void $
execute_
conn
"CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL, age INTEGER NOT NULL)"
void $
execute_
conn
"CREATE TABLE products (id INTEGER PRIMARY KEY, name TEXT NOT NULL, price INTEGER NOT NULL)"
void $
execute_
conn
"CREATE TABLE orders (\
\ id INTEGER PRIMARY KEY, \
\ user_id INTEGER NOT NULL REFERENCES users(id), \
\ product_id INTEGER NOT NULL REFERENCES products(id), \
\ quantity INTEGER NOT NULL)"
seedData :: Connection -> [User] -> [Product] -> [Order] -> IO ()
seedData conn users products orders = runBeamDuckDB conn $ do
runInsert $ insert (_dbUsers testDb) $ insertValues users
runInsert $ insert (_dbProducts testDb) $ insertValues products
runInsert $ insert (_dbOrders testDb) $ insertValues orders
-- Run a test with a fresh in-memory DB populated with the given data
withTestDb ::
[User] ->
[Product] ->
[Order] ->
(Connection -> IO a) ->
IO a
withTestDb users products orders action =
withConnection ":memory:" $ \conn -> do
createTables conn
seedData conn users products orders
action conn