1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18// Tools for dictionaries in IPC context
19
20#ifndef ARROW_IPC_DICTIONARY_H
21#define ARROW_IPC_DICTIONARY_H
22
23#include <cstdint>
24#include <memory>
25#include <unordered_map>
26
27#include "arrow/status.h"
28#include "arrow/util/macros.h"
29#include "arrow/util/visibility.h"
30
31namespace arrow {
32
33class Array;
34class Field;
35
36namespace ipc {
37
38using DictionaryMap = std::unordered_map<int64_t, std::shared_ptr<Array>>;
39using DictionaryTypeMap = std::unordered_map<int64_t, std::shared_ptr<Field>>;
40
41/// \brief Memoization data structure for handling shared dictionaries
42class ARROW_EXPORT DictionaryMemo {
43 public:
44 DictionaryMemo();
45
46 /// \brief Returns KeyError if dictionary not found
47 Status GetDictionary(int64_t id, std::shared_ptr<Array>* dictionary) const;
48
49 /// \brief Return id for dictionary, computing new id if necessary
50 int64_t GetId(const std::shared_ptr<Array>& dictionary);
51
52 /// \brief Return true if dictionary array object is in this memo
53 bool HasDictionary(const std::shared_ptr<Array>& dictionary) const;
54
55 /// \brief Return true if we have a dictionary for the input id
56 bool HasDictionaryId(int64_t id) const;
57
58 /// \brief Add a dictionary to the memo with a particular id. Returns
59 /// KeyError if that dictionary already exists
60 Status AddDictionary(int64_t id, const std::shared_ptr<Array>& dictionary);
61
62 const DictionaryMap& id_to_dictionary() const { return id_to_dictionary_; }
63
64 /// \brief The number of dictionaries stored in the memo
65 int size() const { return static_cast<int>(id_to_dictionary_.size()); }
66
67 private:
68 // Dictionary memory addresses, to track whether a dictionary has been seen
69 // before
70 std::unordered_map<intptr_t, int64_t> dictionary_to_id_;
71
72 // Map of dictionary id to dictionary array
73 DictionaryMap id_to_dictionary_;
74
75 ARROW_DISALLOW_COPY_AND_ASSIGN(DictionaryMemo);
76};
77
78} // namespace ipc
79} // namespace arrow
80
81#endif // ARROW_IPC_DICTIONARY_H
82