{-# LANGUAGE BlockArguments #-} {-# LANGUAGE ForeignFunctionInterface #-} module ArrowInterfaceDeprecatedTests (tests) where import Control.Exception (bracket, finally) import Control.Monad (unless, when) import Data.Int (Int32, Int64) import Data.List (isInfixOf) import Data.Maybe (isNothing) import Database.DuckDB.FFI import Foreign.C.String (peekCString, peekCStringLen, withCString) import Foreign.C.Types (CChar) import Foreign.Marshal.Alloc (alloca) import Foreign.Ptr (FunPtr, Ptr, castPtr, nullFunPtr, nullPtr, plusPtr) import Foreign.Storable (peek, peekElemOff, poke) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) import Utils (withConnection, withDatabase, withResult) tests :: TestTree tests = testGroup "Deprecated Arrow Interface" [ queryArrowExposesSchemaAndArrays , queryArrowReportsErrors , preparedArrowSchemaMatchesStatement , executePreparedArrowProducesRows , resultArrowArrayMirrorsChunk , arrowRowsChangedReflectsMutations , arrowArrayScanRegistersView , arrowStreamScanRegistersView , arrowPointerHelpersClearInternalState ] -- basic query --------------------------------------------------------------- queryArrowExposesSchemaAndArrays :: TestTree queryArrowExposesSchemaAndArrays = testCase "query_arrow exposes schema metadata and arrays" $ withDatabase \db -> withConnection db \conn -> withSuccessfulArrow conn "SELECT 1::INTEGER AS id, 'duck'::VARCHAR AS label" \arrow -> do columnCount <- c_duckdb_arrow_column_count arrow columnCount @?= 2 rowCount <- c_duckdb_arrow_row_count arrow rowCount @?= 1 rowsChanged <- c_duckdb_arrow_rows_changed arrow rowsChanged @?= 0 errPtr <- c_duckdb_query_arrow_error arrow when (errPtr /= nullPtr) $ do errMsg <- peekCString errPtr errMsg @?= "" -- schema/array access validated in dedicated scan test -- error handling ------------------------------------------------------------ queryArrowReportsErrors :: TestTree queryArrowReportsErrors = testCase "query_arrow surfaces execution errors" $ withDatabase \db -> withConnection db \conn -> withCString "SELECT * FROM missing_table" \querySql -> alloca \arrowPtr -> do poke arrowPtr nullPtr state <- c_duckdb_query_arrow conn querySql arrowPtr state @?= DuckDBError arrow <- peek arrowPtr assertBool "arrow result should still be allocated on error" (arrow /= nullPtr) errPtr <- c_duckdb_query_arrow_error arrow assertBool "error message should be present" (errPtr /= nullPtr) errMsg <- peekCString errPtr assertBool "error message should mention missing_table" ("missing_table" `isInfixOf` errMsg) c_duckdb_destroy_arrow arrowPtr -- prepared statements ------------------------------------------------------- preparedArrowSchemaMatchesStatement :: TestTree preparedArrowSchemaMatchesStatement = testCase "prepared_arrow_schema reflects projected columns" $ withDatabase \db -> withConnection db \conn -> withPrepared conn "SELECT id, label FROM (VALUES (1, 'a')) AS t(id, label)" \stmt -> do alloca \schemaOut -> do poke schemaOut nullPtr stSchema <- c_duckdb_prepared_arrow_schema stmt schemaOut case stSchema of DuckDBSuccess -> do schemaWrapper <- peek schemaOut assertBool "prepared arrow schema pointer should not be null" (schemaWrapper /= nullPtr) DuckDBError -> do schemaWrapper <- peek schemaOut schemaWrapper @?= nullPtr errPtr <- c_duckdb_prepare_error stmt when (errPtr /= nullPtr) $ do errMsg <- peekCString errPtr errMsg @?= "" executePreparedArrowProducesRows :: TestTree executePreparedArrowProducesRows = testCase "execute_prepared_arrow materialises a result set" $ withDatabase \db -> withConnection db \conn -> withPrepared conn "SELECT ?::INTEGER + 5 AS computed" \stmt -> do c_duckdb_bind_int32 stmt 1 (5 :: Int32) >>= (@?= DuckDBSuccess) withPreparedArrow stmt \arrow -> do rowCount <- c_duckdb_arrow_row_count arrow rowCount @?= 1 colCount <- c_duckdb_arrow_column_count arrow colCount @?= 1 -- result conversion --------------------------------------------------------- resultArrowArrayMirrorsChunk :: TestTree resultArrowArrayMirrorsChunk = -- NOTE: duckdb_result_arrow_array writes release callbacks into the -- provided ArrowArray storage. Always zero-initialise the struct -- (see zeroArrowArray) and honour the release function before the -- stack memory goes out of scope. testCase "result_arrow_array converts materialised chunks to Arrow arrays" $ withDatabase \db -> withConnection db \conn -> do execStatement conn "CREATE TABLE arrow_chunks(id BIGINT, label VARCHAR);" execStatement conn "INSERT INTO arrow_chunks VALUES (10, 'ten'), (11, 'eleven');" withResult conn "SELECT id, label FROM arrow_chunks ORDER BY id" \resPtr -> do chunk <- c_duckdb_result_get_chunk resPtr 0 assertBool "fetch_chunk returned a null chunk" (chunk /= nullPtr) chunkSize <- c_duckdb_data_chunk_get_size chunk chunkCols <- c_duckdb_data_chunk_get_column_count chunk alloca \arrowArrayPtr -> do poke arrowArrayPtr zeroArrowArray let duckArray :: DuckDBArrowArray duckArray = castPtr arrowArrayPtr alloca \arrayOut -> do poke arrayOut duckArray c_duckdb_result_arrow_array resPtr chunk arrayOut array <- peek arrowArrayPtr arrowArrayLength array @?= fromIntegral chunkSize arrowArrayChildCount array @?= fromIntegral chunkCols validateChunkChildren array releaseArrowArray arrowArrayPtr destroyChunk chunk -- rows changed -------------------------------------------------------------- arrowRowsChangedReflectsMutations :: TestTree arrowRowsChangedReflectsMutations = testCase "arrow_rows_changed reports mutation counts" $ withDatabase \db -> withConnection db \conn -> do execStatement conn "CREATE TABLE arrow_changes(val INTEGER);" withSuccessfulArrow conn "INSERT INTO arrow_changes VALUES (1), (2), (3)" \arrow -> do rowCount <- c_duckdb_arrow_row_count arrow assertBool "modification result should not report negative rows" (rowCount >= 0) changed <- c_duckdb_arrow_rows_changed arrow assertBool "rows_changed should report positive count" (changed > 0) -- arrow scans ---------------------------------------------------------------- arrowArrayScanRegistersView :: TestTree arrowArrayScanRegistersView = -- NOTE: Both the schema and array buffers must be initialised to -- zeroed Arrow structures before calling the query helpers. DuckDB -- fills in release callbacks and expects us to invoke -- duckdb_destroy_arrow_stream on the out stream. testCase "arrow_array_scan registers a view and yields a release stream" $ withDatabase \db -> withConnection db \conn -> do execStatement conn "CREATE TABLE arrow_scan_source(i BIGINT, label VARCHAR);" execStatement conn "INSERT INTO arrow_scan_source VALUES (5, 'five'), (6, 'six');" withSuccessfulArrow conn "SELECT i, label FROM arrow_scan_source ORDER BY i" \arrow -> do alloca \schemaStorage -> do poke schemaStorage zeroArrowSchema let schemaHandle = castPtr schemaStorage :: DuckDBArrowSchema alloca \schemaOut -> do poke schemaOut schemaHandle schemaState <- c_duckdb_query_arrow_schema arrow schemaOut schemaState @?= DuckDBSuccess alloca \arrayStorage -> do poke arrayStorage zeroArrowArray let arrayHandle = castPtr arrayStorage :: DuckDBArrowArray alloca \arrayOut -> do poke arrayOut arrayHandle arrayState <- c_duckdb_query_arrow_array arrow arrayOut arrayState @?= DuckDBSuccess withCString "arrow_array_view" \viewName -> alloca \streamOut -> do poke streamOut nullPtr scanState <- c_duckdb_arrow_array_scan conn viewName schemaHandle arrayHandle streamOut scanState @?= DuckDBSuccess streamWrapper <- peek streamOut assertBool "arrow_array_scan returned a null stream handle" (streamWrapper /= nullPtr) withResult conn "SELECT COUNT(*) FROM arrow_array_view" \resPtr -> do count <- c_duckdb_value_int64 resPtr 0 0 count @?= 2 c_duckdb_destroy_arrow_stream streamOut releaseArrowArray arrayStorage releaseArrowSchema schemaStorage arrowStreamScanRegistersView :: TestTree arrowStreamScanRegistersView = -- NOTE: duckdb_arrow_array_scan returns DuckDBError if the target -- view name already exists (even if it is a table). Use a fresh -- view name that will be dropped implicitly when the connection -- closes. testCase "arrow_scan registers a view from an Arrow stream" $ withDatabase \db -> withConnection db \conn -> do execStatement conn "CREATE TABLE arrow_stream_source(i BIGINT, label VARCHAR);" execStatement conn "INSERT INTO arrow_stream_source VALUES (7, 'seven'), (8, 'eight');" withSuccessfulArrow conn "SELECT i, label FROM arrow_stream_source ORDER BY i" \arrow -> do alloca \schemaStorage -> finally ( do poke schemaStorage zeroArrowSchema let schemaHandle = castPtr schemaStorage :: DuckDBArrowSchema alloca \schemaOut -> do poke schemaOut schemaHandle schemaState <- c_duckdb_query_arrow_schema arrow schemaOut schemaState @?= DuckDBSuccess alloca \arrayStorage -> finally ( do poke arrayStorage zeroArrowArray let arrayHandle = castPtr arrayStorage :: DuckDBArrowArray alloca \arrayOut -> do poke arrayOut arrayHandle arrayState <- c_duckdb_query_arrow_array arrow arrayOut arrayState @?= DuckDBSuccess withCString "arrow_stream_array_view" \sourceView -> alloca \streamOut -> do poke streamOut nullPtr arrayScanState <- c_duckdb_arrow_array_scan conn sourceView schemaHandle arrayHandle streamOut arrayScanState @?= DuckDBSuccess streamHandle <- peek streamOut assertBool "arrow_array_scan returned a null stream" (streamHandle /= nullPtr) withCString "arrow_stream_view" \streamView -> do streamScanState <- c_duckdb_arrow_scan conn streamView streamHandle streamScanState @?= DuckDBSuccess withResult conn "SELECT COUNT(*) FROM arrow_stream_view" \resPtr -> do count <- c_duckdb_value_int64 resPtr 0 0 count @?= 2 c_duckdb_destroy_arrow_stream streamOut ) (releaseArrowArray arrayStorage) ) (releaseArrowSchema schemaStorage) arrowPointerHelpersClearInternalState :: TestTree arrowPointerHelpersClearInternalState = -- NOTE: The deprecated Arrow handles are thin wrappers whose -- internal_ptr fields can be inspected/reset via the exported -- duckdbArrow{Schema,Array,Stream}{Internal,Clear} helpers. We avoid -- direct foreign imports here so future changes to the shims only -- require updates in the Arrow module. testCase "arrow pointer helper functions read and clear internal_ptr fields" $ alloca \schemaFieldPtr -> alloca \arrayFieldPtr -> alloca \streamFieldPtr -> do let schemaField = schemaFieldPtr :: Ptr (Ptr ()) schemaHandle = castPtr schemaField :: DuckDBArrowSchema schemaSentinel = ArrowSchemaPtr (castPtr schemaField) poke schemaField (castPtr (unArrowSchemaPtr schemaSentinel)) schemaPtr <- duckdbArrowSchemaInternal schemaHandle fmap unArrowSchemaPtr schemaPtr @?= Just (unArrowSchemaPtr schemaSentinel) duckdbArrowSchemaClear schemaHandle schemaCleared <- duckdbArrowSchemaInternal schemaHandle assertBool "schema internal_ptr should be cleared" (isNothing schemaCleared) let arrayField = arrayFieldPtr :: Ptr (Ptr ()) arrayHandle = castPtr arrayField :: DuckDBArrowArray arraySentinel = ArrowArrayPtr (castPtr arrayField) poke arrayField (castPtr (unArrowArrayPtr arraySentinel)) arrayPtr <- duckdbArrowArrayInternal arrayHandle fmap unArrowArrayPtr arrayPtr @?= Just (unArrowArrayPtr arraySentinel) duckdbArrowArrayClear arrayHandle arrayCleared <- duckdbArrowArrayInternal arrayHandle assertBool "array internal_ptr should be cleared" (isNothing arrayCleared) let streamField = streamFieldPtr :: Ptr (Ptr ()) streamHandle = castPtr streamField :: DuckDBArrowStream streamSentinel = ArrowStreamPtr (castPtr streamField) poke streamField (castPtr (unArrowStreamPtr streamSentinel)) streamPtr <- duckdbArrowStreamInternal streamHandle fmap unArrowStreamPtr streamPtr @?= Just (unArrowStreamPtr streamSentinel) duckdbArrowStreamClear streamHandle streamCleared <- duckdbArrowStreamInternal streamHandle assertBool "stream internal_ptr should be cleared" (isNothing streamCleared) zeroArrowArray :: ArrowArray zeroArrowArray = ArrowArray { arrowArrayLength = 0 , arrowArrayNullCount = 0 , arrowArrayOffset = 0 , arrowArrayBufferCount = 0 , arrowArrayChildCount = 0 , arrowArrayBuffers = nullPtr , arrowArrayChildren = nullPtr , arrowArrayDictionary = nullPtr , arrowArrayRelease = nullFunPtr , arrowArrayPrivateData = nullPtr } zeroArrowSchema :: ArrowSchema zeroArrowSchema = ArrowSchema { arrowSchemaFormat = nullPtr , arrowSchemaName = nullPtr , arrowSchemaMetadata = nullPtr , arrowSchemaFlags = 0 , arrowSchemaChildCount = 0 , arrowSchemaChildren = nullPtr , arrowSchemaDictionary = nullPtr , arrowSchemaRelease = nullFunPtr , arrowSchemaPrivateData = nullPtr } validateChunkChildren :: ArrowArray -> IO () validateChunkChildren array = do let expectedIds = [10, 11] :: [Int64] expectedLabels = ["ten", "eleven"] childCount = fromIntegral (arrowArrayChildCount array) :: Int rowCount = fromIntegral (arrowArrayLength array) :: Int childCount @?= 2 rowCount @?= length expectedIds let childrenPtr = arrowArrayChildren array assertBool "Arrow array did not expose child arrays" (childrenPtr /= nullPtr) let ensurePtr name ptrPred = unless ptrPred $ assertFailure ("Arrow array " ++ name ++ " pointer is null") intChildPtr <- peekElemOff childrenPtr 0 ensurePtr "integer child" (intChildPtr /= nullPtr) intChild <- peek intChildPtr arrowArrayNullCount intChild @?= 0 let intBufferCount = fromIntegral (arrowArrayBufferCount intChild) :: Int assertBool "Integer child did not expose the expected buffers" (intBufferCount >= 2) let intBuffers = arrowArrayBuffers intChild ensurePtr "integer buffers" (intBuffers /= nullPtr) valueBufferRaw <- peekElemOff intBuffers 1 ensurePtr "integer value buffer" (valueBufferRaw /= nullPtr) let valueBuffer = castPtr valueBufferRaw :: Ptr Int64 values <- mapM (peekElemOff valueBuffer) [0 .. rowCount - 1] values @?= expectedIds strChildPtr <- peekElemOff childrenPtr 1 ensurePtr "varchar child" (strChildPtr /= nullPtr) strChild <- peek strChildPtr arrowArrayNullCount strChild @?= 0 let strBufferCount = fromIntegral (arrowArrayBufferCount strChild) :: Int assertBool "Varchar child did not expose the expected buffers" (strBufferCount >= 3) let strBuffers = arrowArrayBuffers strChild ensurePtr "varchar buffers" (strBuffers /= nullPtr) offsetsRaw <- peekElemOff strBuffers 1 dataRaw <- peekElemOff strBuffers 2 ensurePtr "varchar offsets buffer" (offsetsRaw /= nullPtr) ensurePtr "varchar data buffer" (dataRaw /= nullPtr) let offsetsPtr = castPtr offsetsRaw :: Ptr Int32 dataPtr = castPtr dataRaw :: Ptr CChar labels <- mapM ( \idx -> do start <- fromIntegral <$> peekElemOff offsetsPtr idx end <- fromIntegral <$> peekElemOff offsetsPtr (idx + 1) peekCStringLen (dataPtr `plusPtr` start, end - start) ) [0 .. rowCount - 1] labels @?= expectedLabels releaseArrowArray :: Ptr ArrowArray -> IO () releaseArrowArray arrayPtr = do array <- peek arrayPtr let releaseFun = arrowArrayRelease array when (releaseFun /= nullFunPtr) $ do let release = mkArrowArrayRelease releaseFun release arrayPtr releaseArrowSchema :: Ptr ArrowSchema -> IO () releaseArrowSchema schemaPtr = do schema <- peek schemaPtr let releaseFun = arrowSchemaRelease schema when (releaseFun /= nullFunPtr) $ do let release = mkArrowSchemaRelease releaseFun release schemaPtr foreign import ccall "dynamic" mkArrowArrayRelease :: FunPtr (Ptr ArrowArray -> IO ()) -> Ptr ArrowArray -> IO () foreign import ccall "dynamic" mkArrowSchemaRelease :: FunPtr (Ptr ArrowSchema -> IO ()) -> Ptr ArrowSchema -> IO () withSuccessfulArrow :: DuckDBConnection -> String -> (DuckDBArrow -> IO a) -> IO a withSuccessfulArrow conn sql action = withCString sql \sqlPtr -> alloca \arrowPtr -> bracket ( do poke arrowPtr nullPtr state <- c_duckdb_query_arrow conn sqlPtr arrowPtr state @?= DuckDBSuccess arrow <- peek arrowPtr assertBool "duckdb_query_arrow returned null result" (arrow /= nullPtr) pure arrow ) (\_ -> c_duckdb_destroy_arrow arrowPtr) action withPrepared :: DuckDBConnection -> String -> (DuckDBPreparedStatement -> IO a) -> IO a withPrepared conn sql action = withCString sql \sqlPtr -> alloca \stmtPtr -> bracket ( do state <- c_duckdb_prepare conn sqlPtr stmtPtr state @?= DuckDBSuccess stmt <- peek stmtPtr assertBool "prepare should produce a statement" (stmt /= nullPtr) pure stmt ) (\_ -> c_duckdb_destroy_prepare stmtPtr) action withPreparedArrow :: DuckDBPreparedStatement -> (DuckDBArrow -> IO a) -> IO a withPreparedArrow stmt action = alloca \arrowPtr -> bracket ( do poke arrowPtr nullPtr state <- c_duckdb_execute_prepared_arrow stmt arrowPtr state @?= DuckDBSuccess arrow <- peek arrowPtr assertBool "execute_prepared_arrow returned null result" (arrow /= nullPtr) pure arrow ) (\_ -> c_duckdb_destroy_arrow arrowPtr) action destroyChunk :: DuckDBDataChunk -> IO () destroyChunk chunk = alloca \ptr -> poke ptr chunk >> c_duckdb_destroy_data_chunk ptr execStatement :: DuckDBConnection -> String -> IO () execStatement conn sql = withCString sql \sqlPtr -> alloca \resPtr -> do st <- c_duckdb_query conn sqlPtr resPtr st @?= DuckDBSuccess c_duckdb_destroy_result resPtr