1//===----------------------------------------------------------------------===//
2// DuckDB
3//
4// duckdb/common/arrow/arrow_wrapper.hpp
5//
6//
7//===----------------------------------------------------------------------===//
8
9#pragma once
10#include "duckdb/common/arrow/arrow.hpp"
11#include "duckdb/common/helper.hpp"
12#include "duckdb/common/preserved_error.hpp"
13
14//! Here we have the internal duckdb classes that interact with Arrow's Internal Header (i.e., duckdb/commons/arrow.hpp)
15namespace duckdb {
16class QueryResult;
17class DataChunk;
18
19class ArrowSchemaWrapper {
20public:
21 ArrowSchema arrow_schema;
22
23 ArrowSchemaWrapper() {
24 arrow_schema.release = nullptr;
25 }
26
27 ~ArrowSchemaWrapper();
28};
29class ArrowArrayWrapper {
30public:
31 ArrowArray arrow_array;
32 ArrowArrayWrapper() {
33 arrow_array.length = 0;
34 arrow_array.release = nullptr;
35 }
36 ~ArrowArrayWrapper();
37};
38
39class ArrowArrayStreamWrapper {
40public:
41 ArrowArrayStream arrow_array_stream;
42 int64_t number_of_rows;
43
44public:
45 void GetSchema(ArrowSchemaWrapper &schema);
46
47 shared_ptr<ArrowArrayWrapper> GetNextChunk();
48
49 const char *GetError();
50
51 ~ArrowArrayStreamWrapper();
52 ArrowArrayStreamWrapper() {
53 arrow_array_stream.release = nullptr;
54 }
55};
56
57class ArrowUtil {
58public:
59 static bool TryFetchChunk(QueryResult *result, idx_t chunk_size, ArrowArray *out, idx_t &result_count,
60 PreservedError &error);
61 static idx_t FetchChunk(QueryResult *result, idx_t chunk_size, ArrowArray *out);
62
63private:
64 static bool TryFetchNext(QueryResult &result, unique_ptr<DataChunk> &out, PreservedError &error);
65};
66} // namespace duckdb
67