forked from oceanbase/seekdb
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathob_arrow_basic.h
More file actions
180 lines (147 loc) · 5.33 KB
/
Copy pathob_arrow_basic.h
File metadata and controls
180 lines (147 loc) · 5.33 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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
/*
* Copyright (c) 2025 OceanBase.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef OB_PARQUET_BASIC_H
#define OB_PARQUET_BASIC_H
#include <arrow/api.h>
#include <arrow/io/api.h>
#include <parquet/arrow/reader.h>
#include <parquet/arrow/writer.h>
#include <parquet/exception.h>
#include <orc/OrcFile.hh>
#include "share/ob_device_manager.h"
#include "sql/engine/table/ob_external_table_access_service.h"
#include "sql/engine/basic/ob_select_into_basic.h"
#include "sql/engine/table/ob_file_prefetch_buffer.h"
namespace oceanbase
{
namespace sql
{
class ObOrcMemPool : public orc::MemoryPool {
public:
void init(uint64_t tenant_id);
virtual char* malloc(uint64_t size) override;
virtual void free(char* p) override;
private:
common::ObMemAttr mem_attr_;
};
class ObOrcOutputStream : public orc::OutputStream {
public:
ObOrcOutputStream(ObFileAppender *file_appender,
ObStorageAppender *storage_appender,
IntoFileLocation file_location,
const ObString &url)
: file_appender_(file_appender),
storage_appender_(storage_appender),
file_location_(file_location),
url_(url.ptr()),
pos_(0)
{}
~ObOrcOutputStream() {}
virtual void write(const void *buf, size_t length) override;
virtual uint64_t getLength() const { return pos_; }
virtual uint64_t getNaturalWriteSize() const { return 1024; }
virtual const std::string& getName() const override {
return url_;
}
virtual void close() {
}
private:
ObFileAppender *file_appender_;
ObStorageAppender *storage_appender_;
IntoFileLocation file_location_;
std::string url_;
int64_t pos_;
};
class ObArrowMemPool : public ::arrow::MemoryPool
{
public:
ObArrowMemPool() : total_alloc_size_(0) {}
void init(uint64_t tenant_id);
virtual arrow::Status Allocate(int64_t size, uint8_t** out) override;
virtual arrow::Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override;
virtual void Free(uint8_t* buffer, int64_t size) override;
virtual void ReleaseUnused() override;
virtual int64_t bytes_allocated() const override;
virtual int64_t max_memory() const override { return -1; }
virtual std::string backend_name() const override { return "Arrow"; }
private:
common::ObArenaAllocator alloc_;
common::ObMemAttr mem_attr_;
arrow::internal::MemoryPoolStats stats_;
int64_t total_alloc_size_;
};
class ObArrowFile : public arrow::io::RandomAccessFile {
public:
ObArrowFile(ObExternalDataAccessDriver &file_reader, const char *file_name,
arrow::MemoryPool *pool, ObFilePrefetchBuffer &file_prefetch_buffer) :
file_reader_(file_reader),
file_name_(file_name), pool_(pool), file_prefetch_buffer_(file_prefetch_buffer)
{
file_prefetch_buffer_.clear();
}
~ObArrowFile() override {
file_reader_.close();
}
int open();
virtual arrow::Status Close() override;
virtual bool closed() const override;
virtual arrow::Result<int64_t> Read(int64_t nbytes, void* out) override;
virtual arrow::Result<std::shared_ptr<arrow::Buffer>> Read(int64_t nbytes) override;
virtual arrow::Result<int64_t> ReadAt(int64_t position, int64_t nbytes, void* out) override;
virtual arrow::Result<std::shared_ptr<arrow::Buffer>> ReadAt(int64_t position, int64_t nbytes) override;
virtual arrow::Status Seek(int64_t position) override;
virtual arrow::Result<int64_t> Tell() const override;
virtual arrow::Result<int64_t> GetSize() override;
private:
ObExternalDataAccessDriver &file_reader_;
const char* file_name_;
arrow::MemoryPool *pool_;
int64_t position_;
ObFilePrefetchBuffer &file_prefetch_buffer_;
};
class ObParquetOutputStream : public arrow::io::OutputStream
{
public:
ObParquetOutputStream (ObFileAppender *file_appender,
ObStorageAppender *storage_appender,
IntoFileLocation file_location,
const ObString &url)
: file_appender_(file_appender),
storage_appender_(storage_appender),
file_location_(file_location),
url_(url),
position_(0)
{}
~ObParquetOutputStream() override {}
// Write methods
// Virtual methods in `arrow::io::Writable`
virtual arrow::Status Write(const void* data, int64_t nbytes) override;
// virtual arrow::Status Write(const std::shared_ptr<arrow::Buffer>& data) override;
// virtual arrow::Status Flush() override;
// Virtual methods in `arrow::io::FileInterface`
virtual arrow::Status Close() override;
virtual bool closed() const override;
virtual arrow::Result<int64_t> Tell() const override;
private:
ObFileAppender *file_appender_;
ObStorageAppender *storage_appender_;
IntoFileLocation file_location_;
const ObString &url_;
int64_t position_;
};
} // end of sql namespace
} // end of oceanbase namespace
#endif // OB_PARQUET_BASIC_H