duckdb-ffi-0.1.4.1: test/VectorTest.hs
{-# LANGUAGE BlockArguments #-}
module VectorTest (tests) where
import Control.Monad (forM_, void, (>=>))
import Data.Bits (setBit)
import Data.Int (Int32)
import Data.Word (Word64)
import Database.DuckDB.FFI
import Foreign.C.String (withCString)
import Foreign.Marshal.Array (allocaArray, withArray)
import Foreign.Ptr (Ptr, castPtr, nullPtr, plusPtr)
import Foreign.Storable (peek, peekElemOff, poke, pokeElemOff, sizeOf)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, testCase, (@?=))
import Utils (clearValidityBit, setAllValid, withLogicalType, withSelectionVector, withValue, withVectorOfType)
-- | Entry point for vector-centric tests.
tests :: TestTree
tests =
testGroup
"Vector Interface"
[ vectorDataAccess
, vectorValidityMask
, listVectorChildManagement
, arrayVectorChildAccess
, vectorSliceWithSelection
, structVectorChildAccess
, vectorReferenceValue
]
-- | Create a numeric vector, write elements, and observe the stored values.
vectorDataAccess :: TestTree
vectorDataAccess =
testCase "write and read integer vector data" $ do
withLogicalType (c_duckdb_create_logical_type DuckDBTypeInteger) \intType ->
withVectorOfType intType 4 \vec -> do
colType <- c_duckdb_vector_get_column_type vec
withLogicalType (pure colType) (c_duckdb_get_type_id >=> (@?= DuckDBTypeInteger))
rawPtr <- c_duckdb_vector_get_data vec
let dataPtr = castPtr rawPtr :: Ptr Int32
forM_ (zip [0 ..] [10, 20, 30, 40]) \(idx, val) -> poke (dataPtr `plusElem` idx) val
forM_ (zip [0 ..] [10, 20, 30, 40]) \(idx, val) -> peek (dataPtr `plusElem` idx) >>= (@?= val)
-- | Ensure validity mask allocation and manipulate individual bits.
vectorValidityMask :: TestTree
vectorValidityMask =
testCase "ensure validity mask and clear single entry" $ do
withLogicalType (c_duckdb_create_logical_type DuckDBTypeInteger) \intType ->
withVectorOfType intType 4 \vec -> do
void (c_duckdb_vector_ensure_validity_writable vec)
validity <- c_duckdb_vector_get_validity vec
assertBool "validity mask should exist" (validity /= nullPtr)
setAllValid validity 4
clearValidityBit validity 2
current <- peek validity
let expected = foldl setBit (0 :: Word64) [0, 1, 3]
current @?= expected
-- | Exercise list vector child management APIs.
listVectorChildManagement :: TestTree
listVectorChildManagement =
testCase "list vector reserves space and reports size" $ do
withLogicalType (c_duckdb_create_logical_type DuckDBTypeInteger) \intType ->
withLogicalType (c_duckdb_create_list_type intType) \listType ->
withVectorOfType listType 2 \listVec -> do
childVec0 <- c_duckdb_list_vector_get_child listVec
assertBool "list child vector should be non-null" (childVec0 /= nullPtr)
initialSize <- c_duckdb_list_vector_get_size listVec
initialSize @?= 0
reserveState <- c_duckdb_list_vector_reserve listVec 5
reserveState @?= DuckDBSuccess
sizeState <- c_duckdb_list_vector_set_size listVec 5
sizeState @?= DuckDBSuccess
c_duckdb_list_vector_get_size listVec >>= (@?= 5)
childVec <- c_duckdb_list_vector_get_child listVec
metaRaw <- c_duckdb_vector_get_data listVec
let metaWords = castPtr metaRaw :: Ptr Word64
-- row 0 -> offset 0 length 3, row 1 -> offset 3 length 2
pokeElemOff metaWords 0 0
pokeElemOff metaWords 1 3
pokeElemOff metaWords 2 3
pokeElemOff metaWords 3 2
childRaw <- c_duckdb_vector_get_data childVec
let childInts = castPtr childRaw :: Ptr Int32
payload = [11, 12, 13, 21, 22 :: Int32]
forM_ (zip [0 ..] payload) (uncurry (pokeElemOff childInts))
fetched <- mapM (peekElemOff childInts) [0 .. length payload - 1]
fetched @?= payload
peekElemOff metaWords 1 >>= (@?= 3)
peekElemOff metaWords 3 >>= (@?= 2)
-- | Verify array child vector exposes a flat buffer sized by row * array length.
arrayVectorChildAccess :: TestTree
arrayVectorChildAccess =
testCase "array vector child flattens elements" $ do
withLogicalType (c_duckdb_create_logical_type DuckDBTypeInteger) \intType ->
withLogicalType (c_duckdb_create_array_type intType 3) \arrayType ->
withVectorOfType arrayType 2 \arrayVec -> do
childVec <- c_duckdb_array_vector_get_child arrayVec
childType <- c_duckdb_vector_get_column_type childVec
withLogicalType (pure childType) (c_duckdb_get_type_id >=> (@?= DuckDBTypeInteger))
childRaw <- c_duckdb_vector_get_data childVec
let childInts = castPtr childRaw :: Ptr Int32
payload = [1, 2, 3, 4, 5, 6 :: Int32]
forM_ (zip [0 ..] payload) (uncurry (pokeElemOff childInts))
mapM (peekElemOff childInts) [0 .. length payload - 1] >>= (@?= payload)
-- | Slice a vector with a selection vector and materialize the dictionary.
vectorSliceWithSelection :: TestTree
vectorSliceWithSelection =
testCase "slice vector materializes dictionary order" $ do
withLogicalType (c_duckdb_create_logical_type DuckDBTypeInteger) \intType ->
withVectorOfType intType 4 \vec -> do
dataRaw <- c_duckdb_vector_get_data vec
let vecData = castPtr dataRaw :: Ptr Int32
forM_ (zip [0 ..] [10, 20, 30, 40 :: Int32]) (uncurry (pokeElemOff vecData))
withSelectionVector 2 \sel -> do
selPtr <- c_duckdb_selection_vector_get_data_ptr sel
pokeElemOff selPtr 0 (fromIntegral (3 :: Int))
pokeElemOff selPtr 1 (fromIntegral (1 :: Int))
c_duckdb_slice_vector vec sel 2
withVectorOfType intType 2 \materialized ->
withSelectionVector 2 \copySel -> do
copyPtr <- c_duckdb_selection_vector_get_data_ptr copySel
pokeElemOff copyPtr 0 0
pokeElemOff copyPtr 1 1
c_duckdb_vector_copy_sel vec materialized copySel 2 0 0
matRaw <- c_duckdb_vector_get_data materialized
let matPtr = castPtr matRaw :: Ptr Int32
mapM (peekElemOff matPtr) [0, 1] >>= (@?= [40, 20])
-- | Ensure struct vectors expose individual child vectors with matching types.
structVectorChildAccess :: TestTree
structVectorChildAccess =
testCase "struct vector exposes typed child vectors" $ do
withLogicalType (c_duckdb_create_logical_type DuckDBTypeInteger) \intType ->
withLogicalType (c_duckdb_create_logical_type DuckDBTypeDouble) \doubleType ->
allocaArray 2 \typesPtr -> do
pokeElemOff typesPtr 0 intType
pokeElemOff typesPtr 1 doubleType
withCString "ints" \name0 ->
withCString "doubles" \name1 ->
withArray [name0, name1] \namesPtr ->
withLogicalType (c_duckdb_create_struct_type typesPtr namesPtr 2) \structType ->
withVectorOfType structType 1 \structVec -> do
intChild <- c_duckdb_struct_vector_get_child structVec 0
dblChild <- c_duckdb_struct_vector_get_child structVec 1
intChildType <- c_duckdb_vector_get_column_type intChild
withLogicalType (pure intChildType) (c_duckdb_get_type_id >=> (@?= DuckDBTypeInteger))
dblChildType <- c_duckdb_vector_get_column_type dblChild
withLogicalType (pure dblChildType) (c_duckdb_get_type_id >=> (@?= DuckDBTypeDouble))
intRaw <- c_duckdb_vector_get_data intChild
dblRaw <- c_duckdb_vector_get_data dblChild
let intPtr = castPtr intRaw :: Ptr Int32
dblPtr = castPtr dblRaw :: Ptr Double
pokeElemOff intPtr 0 7
pokeElemOff dblPtr 0 3.5
peekElemOff intPtr 0 >>= (@?= 7)
peekElemOff dblPtr 0 >>= (@?= 3.5)
-- | Copy a scalar duckdb_value into a vector without materializing a chunk.
vectorReferenceValue :: TestTree
vectorReferenceValue =
testCase "vector_reference_value writes scalar contents" $ do
withLogicalType (c_duckdb_create_logical_type DuckDBTypeInteger) \intType ->
withVectorOfType intType 1 \vec ->
withValue (c_duckdb_create_int32 123) \value -> do
c_duckdb_vector_reference_value vec value
raw <- c_duckdb_vector_get_data vec
let dataPtr = castPtr raw :: Ptr Int32
peekElemOff dataPtr 0 >>= (@?= 123)
-- helpers -------------------------------------------------------------------
plusElem :: Ptr Int32 -> Int -> Ptr Int32
plusElem base idx = base `plusPtr` (idx * sizeOf (undefined :: Int32))