paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1543526576
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2151,53 +2188,96 @@ class ExportedArrayStream {
int64_t next_batch_num() { return private_data()->batch_num_++; }
- struct ArrowArrayStream* stream_;
+ StreamType* stream_;
};
} // namespace
Status ExportRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<RecordBatchReader>::Make(std::move(reader), out);
+ return ExportedArrayStream<RecordBatchReader, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(reader), out);
}
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<ChunkedArray>::Make(std::move(chunked_array),
out);
+ return ExportedArrayStream<ChunkedArray, struct ArrowArrayStream,
+ struct
ArrowArray>::Make(std::move(chunked_array), out);
+}
+
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(reader->device_type());
+ return ExportedArrayStream<RecordBatchReader, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(reader),
out);
+}
+
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(device_type);
+ return ExportedArrayStream<ChunkedArray, struct ArrowDeviceArrayStream,
+ struct
ArrowDeviceArray>::Make(std::move(chunked_array),
+ out);
}
//////////////////////////////////////////////////////////////////////////
// C stream import
namespace {
+template <typename StreamTraits, typename ArrayTraits>
class ArrayStreamReader {
+ protected:
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamReader(struct ArrowArrayStream* stream) {
- ArrowArrayStreamMove(stream, &stream_);
- DCHECK(!ArrowArrayStreamIsReleased(&stream_));
+ explicit ArrayStreamReader(StreamType* stream,
+ const DeviceMemoryMapper& mapper =
DefaultDeviceMapper)
+ : mapper_{mapper} {
+ StreamTraits::MoveFunc(stream, &stream_);
+ DCHECK(!StreamTraits::IsReleasedFunc(&stream_));
}
~ArrayStreamReader() { ReleaseStream(); }
void ReleaseStream() {
- if (!ArrowArrayStreamIsReleased(&stream_)) {
- ArrowArrayStreamRelease(&stream_);
- }
- DCHECK(ArrowArrayStreamIsReleased(&stream_));
+ // all our trait release funcs check IsReleased so we don't
+ // need to repeat it here
+ StreamTraits::ReleaseFunc(&stream_);
+ DCHECK(StreamTraits::IsReleasedFunc(&stream_));
}
protected:
- Status ReadNextArrayInternal(struct ArrowArray* array) {
- ArrowArrayMarkReleased(array);
+ Status ReadNextArrayInternal(ArrayType* array) {
+ ArrayTraits::MarkReleased(array);
Status status = StatusFromCError(stream_.get_next(&stream_, array));
- if (!status.ok() && !ArrowArrayIsReleased(array)) {
- ArrowArrayRelease(array);
+ if (!status.ok() && !ArrayTraits::IsReleasedFunc(array)) {
+ ArrayTraits::ReleaseFunc(array);
}
return status;
}
+ Result<std::shared_ptr<RecordBatch>> ImportRecordBatchInternal(
+ ArrayType* array, std::shared_ptr<Schema> schema) {
Review Comment:
Your call! That was feedback I got the last time I tried to use `if
constexpr` in this section of code 🙂
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
