-
Notifications
You must be signed in to change notification settings - Fork 191
Expand file tree
/
Copy pathscan.cpp
More file actions
98 lines (83 loc) · 2.58 KB
/
Copy pathscan.cpp
File metadata and controls
98 lines (83 loc) · 2.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
#include "vortex/array.hpp"
#include "vortex/common.hpp"
#include "vortex/dtype.hpp"
#include "vortex/error.hpp"
#include "vortex/scan.hpp"
#include <vortex.h>
#include <mutex>
#include <optional>
namespace vortex {
using detail::Access;
using detail::throw_on_error;
ArrowStream::ArrowStream(ArrowStream &&other) noexcept
: session_(std::move(other.session_)), stream_(other.stream_) {
other.stream_ = ArrowArrayStream {};
}
ArrowStream &ArrowStream::operator=(ArrowStream &&other) noexcept {
if (this != &other) {
if (stream_.release != nullptr) {
stream_.release(&stream_);
}
session_ = std::move(other.session_);
stream_ = other.stream_;
other.stream_ = ArrowArrayStream {};
}
return *this;
}
ArrowStream::~ArrowStream() {
if (stream_.release != nullptr) {
stream_.release(&stream_);
}
}
void Partition::Deleter::operator()(vx_partition *ptr) const noexcept {
vx_partition_free(ptr);
}
Estimate Partition::row_count() const {
vx_estimate raw {};
vx_error *error = nullptr;
vx_partition_row_count(handle_.get(), &raw, &error);
throw_on_error(error);
return Estimate(raw);
}
std::optional<Array> Partition::next() {
vx_error *error = nullptr;
const vx_array *array = vx_partition_next(handle_.get(), &error);
throw_on_error(error);
if (array == nullptr) {
return std::nullopt;
}
return Access::adopt<Array>(array);
}
ArrowStream Partition::into_arrow_stream() && {
ArrowArrayStream stream {};
vx_error *error = nullptr;
// Consumes handle even on error
vx_partition_scan_arrow(Access::c_ptr(session_), handle_.release(), &stream, &error);
throw_on_error(error);
return Access::adopt<ArrowStream>(std::move(session_), stream);
}
void Scan::Deleter::operator()(vx_scan *ptr) const noexcept {
vx_scan_free(ptr);
}
DataType Scan::dtype() const {
vx_error *error = nullptr;
const vx_dtype *out = vx_scan_dtype(handle_.get(), &error);
throw_on_error(error);
return Access::adopt<DataType>(out);
}
std::optional<Partition> Scan::next_partition() {
vx_partition *partition = nullptr;
{
const std::lock_guard<std::mutex> guard(*mutex_);
vx_error *error = nullptr;
partition = vx_scan_next_partition(handle_.get(), &error);
throw_on_error(error);
}
if (partition == nullptr) {
return std::nullopt;
}
return Access::adopt<Partition>(partition, session_);
}
} // namespace vortex