Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions adbc/src/connector/adbc_connector.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,21 @@ static std::vector<std::string> splitCommaSeparated(const std::string& input) {
return result;
}

static std::vector<std::optional<common::ArrowLogicalTypeInfo>> resolveColumnLogicalTypeInfos(
const ArrowSchemaWrapper& schema) {
std::vector<std::optional<common::ArrowLogicalTypeInfo>> result;
if (schema.n_children <= 0 || schema.children == nullptr) {
return result;
}
result.reserve(static_cast<size_t>(schema.n_children));
for (int64_t i = 0; i < schema.n_children; ++i) {
result.push_back(schema.children[i] == nullptr ?
std::nullopt :
common::tryGetArrowLogicalTypeInfo(schema.children[i]));
}
return result;
}

ADBCQueryResult::~ADBCQueryResult() {
if (stream.release) {
stream.release(&stream);
Expand Down Expand Up @@ -200,6 +215,7 @@ std::unique_ptr<ADBCQueryResult> ADBCConnector::executeQuery(const std::string&
throw common::RuntimeException{
std::format("ArrowArrayStream.get_schema failed: {}", streamError)};
}
result->logicalTypeInfos = resolveColumnLogicalTypeInfos(result->schema);
// Fully consume the ADBC stream before returning control to the execution engine. Some drivers
// keep query state pending while the stream is open, and later catalog calls can close it.
while (true) {
Expand Down
5 changes: 4 additions & 1 deletion adbc/src/function/adbc_scan.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -102,8 +102,11 @@ common::offset_t ADBCScanFunction::tableFunc(const function::TableFuncInput& inp
auto srcOffset = localState->array.children[i]->offset + localState->offset;
common::ArrowNullMaskTree mask(schema.children[i], localState->array.children[i], srcOffset,
count);
const auto* logicalTypeInfo = i < sharedState->queryResult->logicalTypeInfos.size() ?
&sharedState->queryResult->logicalTypeInfos[i] :
nullptr;
common::ArrowConverter::fromArrowArray(schema.children[i], localState->array.children[i],
output.dataChunk.getValueVectorMutable(i), &mask, srcOffset, 0, count);
output.dataChunk.getValueVectorMutable(i), &mask, srcOffset, 0, count, logicalTypeInfo);
}
localState->offset += count;
return count;
Expand Down
2 changes: 2 additions & 0 deletions adbc/src/include/connector/adbc_connector.h
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

#include "binder/bound_attach_info.h"
#include "common/arrow/arrow.h"
#include "common/arrow/arrow_schema_metadata.h"
#include "common/types/types.h"

#if __has_include(<arrow-adbc/adbc.h>) && __has_include(<arrow-adbc/adbc_driver_manager.h>)
Expand Down Expand Up @@ -33,6 +34,7 @@ struct ADBCQueryResult {
AdbcStatement statement{};
ArrowArrayStream stream{};
ArrowSchemaWrapper schema{};
std::vector<std::optional<common::ArrowLogicalTypeInfo>> logicalTypeInfos;
std::vector<ArrowArrayWrapper> arrays;
uint64_t nextArrayIdx = 0;
bool connectionInitialized = false;
Expand Down
Loading