duckdb::ResultArrowArrayStreamWrapper Class Reference

Collaboration diagram for duckdb::ResultArrowArrayStreamWrapper:

Public Member Functions
	ResultArrowArrayStreamWrapper (unique_ptr< QueryResult > result, idx_t batch_size)

Public Attributes
ArrowArrayStream	stream
unique_ptr< QueryResult >	result
ErrorData	last_error
idx_t	batch_size
vector< LogicalType >	column_types
vector< string >	column_names
unique_ptr< ChunkScanState >	scan_state
unordered_map< idx_t, const shared_ptr< ArrowTypeExtensionData > >	extension_types

Static Private Member Functions
static int	MyStreamGetSchema (struct ArrowArrayStream *stream, struct ArrowSchema *out)
static int	MyStreamGetNext (struct ArrowArrayStream *stream, struct ArrowArray *out)
static void	MyStreamRelease (struct ArrowArrayStream *stream)
static const char *	MyStreamGetLastError (struct ArrowArrayStream *stream)

Constructor & Destructor Documentation

ResultArrowArrayStreamWrapper()

duckdb::ResultArrowArrayStreamWrapper::ResultArrowArrayStreamWrapper ( unique_ptr< QueryResult >  result, idx_t  batch_size  )

explicit

We first initialize the private data of the stream

Ceil Approx_Batch_Size/STANDARD_VECTOR_SIZE

We initialize the stream functions

    : result(std::move(result_p)), scan_state(make_uniq<QueryResultChunkScanState>(*result)) {
    stream.private_data = this;
    if (batch_size_p == 0) {
        throw std::runtime_error("Approximate Batch Size of Record Batch MUST be higher than 0");
    }
    batch_size = batch_size_p;
    stream.get_schema = ResultArrowArrayStreamWrapper::MyStreamGetSchema;
    stream.get_next = ResultArrowArrayStreamWrapper::MyStreamGetNext;
    stream.release = ResultArrowArrayStreamWrapper::MyStreamRelease;
    stream.get_last_error = ResultArrowArrayStreamWrapper::MyStreamGetLastError;
 
    extension_types =
        ArrowTypeExtensionData::GetExtensionTypes(*result->client_properties.client_context, result->types);
}

Here is the call graph for this function:

Member Function Documentation

MyStreamGetSchema()

int duckdb::ResultArrowArrayStreamWrapper::MyStreamGetSchema ( struct ArrowArrayStream *  stream, struct ArrowSchema *  out  )

staticprivate

                                                                                                             {
    if (!stream->release) {
        return -1;
    }
    out->release = nullptr;
    auto my_stream = reinterpret_cast<ResultArrowArrayStreamWrapper *>(stream->private_data);
    if (!my_stream->column_types.empty()) {
        try {
            ArrowConverter::ToArrowSchema(out, my_stream->column_types, my_stream->column_names,
                                          my_stream->result->client_properties);
        } catch (std::runtime_error &e) {
            my_stream->last_error = ErrorData(e);
            return -1;
        }
        return 0;
    }
 
    auto &result = *my_stream->result;
    if (result.HasError()) {
        my_stream->last_error = result.GetErrorObject();
        return -1;
    }
    if (result.type == QueryResultType::STREAM_RESULT) {
        auto &stream_result = result.Cast<StreamQueryResult>();
        if (!stream_result.IsOpen()) {
            my_stream->last_error = ErrorData("Query Stream is closed");
            return -1;
        }
    }
    if (my_stream->column_types.empty()) {
        my_stream->column_types = result.types;
        my_stream->column_names = result.names;
    }
    try {
        ArrowConverter::ToArrowSchema(out, my_stream->column_types, my_stream->column_names,
                                      my_stream->result->client_properties);
    } catch (std::runtime_error &e) {
        my_stream->last_error = ErrorData(e);
        return -1;
    }
    return 0;
}

MyStreamGetNext()

int duckdb::ResultArrowArrayStreamWrapper::MyStreamGetNext ( struct ArrowArrayStream *  stream, struct ArrowArray *  out  )

staticprivate

                                                                                                          {
    if (!stream->release) {
        return -1;
    }
    auto my_stream = reinterpret_cast<ResultArrowArrayStreamWrapper *>(stream->private_data);
    auto &result = *my_stream->result;
    auto &scan_state = *my_stream->scan_state;
    if (result.HasError()) {
        my_stream->last_error = result.GetErrorObject();
        return -1;
    }
    if (result.type == QueryResultType::STREAM_RESULT) {
        auto &stream_result = result.Cast<StreamQueryResult>();
        if (!stream_result.IsOpen()) {
            // Nothing to output
            out->release = nullptr;
            return 0;
        }
    }
    if (my_stream->column_types.empty()) {
        my_stream->column_types = result.types;
        my_stream->column_names = result.names;
    }
 
    try {
        idx_t result_count;
        ErrorData error;
        if (!ArrowUtil::TryFetchChunk(scan_state, result.client_properties, my_stream->batch_size, out, result_count,
                                      error, my_stream->extension_types)) {
            D_ASSERT(error.HasError());
            my_stream->last_error = error;
            return -1;
        }
        if (result_count == 0) {
            // Nothing to output
            out->release = nullptr;
        }
    } catch (std::exception &e) {
        my_stream->last_error = ErrorData(e);
        return -1;
    }
 
    return 0;
}

MyStreamRelease()

void duckdb::ResultArrowArrayStreamWrapper::MyStreamRelease ( struct ArrowArrayStream *  stream )

staticprivate

                                                                                   {
    if (!stream || !stream->release) {
        return;
    }
    stream->release = nullptr;
    delete reinterpret_cast<ResultArrowArrayStreamWrapper *>(stream->private_data);
}

MyStreamGetLastError()

const char * duckdb::ResultArrowArrayStreamWrapper::MyStreamGetLastError ( struct ArrowArrayStream *  stream )

staticprivate

                                                                                               {
    if (!stream->release) {
        return "stream was released";
    }
    D_ASSERT(stream->private_data);
    auto my_stream = reinterpret_cast<ResultArrowArrayStreamWrapper *>(stream->private_data);
    return my_stream->last_error.Message().c_str();
}

---

The documentation for this class was generated from the following file:

• external/duckdb/duckdb.cpp