mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-01-03 09:22:30 +08:00
48 lines
1.8 KiB
C++
48 lines
1.8 KiB
C++
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
//
|
|
// 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
|
|
|
|
#pragma once
|
|
|
|
#include <memory>
|
|
#include <utility>
|
|
|
|
#include "common/Channel.h"
|
|
#include "parquet/arrow/reader.h"
|
|
|
|
namespace milvus {
|
|
|
|
struct ArrowTableInfo {
|
|
size_t file_index; // file index in the file list
|
|
size_t row_group_index; // row group index in the file, start from 0
|
|
std::shared_ptr<arrow::Table>
|
|
table; // the arrow table of the row group at row group index in the file index file
|
|
};
|
|
|
|
struct ArrowDataWrapper {
|
|
ArrowDataWrapper() = default;
|
|
ArrowDataWrapper(std::shared_ptr<arrow::RecordBatchReader> reader,
|
|
std::shared_ptr<parquet::arrow::FileReader> arrow_reader,
|
|
std::shared_ptr<uint8_t[]> file_data)
|
|
: reader(std::move(reader)),
|
|
arrow_reader(std::move(arrow_reader)),
|
|
file_data(std::move(file_data)) {
|
|
}
|
|
std::shared_ptr<arrow::RecordBatchReader> reader;
|
|
// file reader must outlive the record batch reader
|
|
std::shared_ptr<parquet::arrow::FileReader> arrow_reader;
|
|
// underlying file data memory, must outlive the arrow reader
|
|
std::shared_ptr<uint8_t[]> file_data;
|
|
std::vector<ArrowTableInfo> arrow_tables;
|
|
};
|
|
using ArrowReaderChannel = Channel<std::shared_ptr<milvus::ArrowDataWrapper>>;
|
|
|
|
} // namespace milvus
|