]>
Commit | Line | Data |
---|---|---|
1 | /* | |
2 | * Licensed to the Apache Software Foundation (ASF) under one or more | |
3 | * contributor license agreements. See the NOTICE file distributed with | |
4 | * this work for additional information regarding copyright ownership. | |
5 | * The ASF licenses this file to You under the Apache License, Version 2.0 | |
6 | * (the "License"); you may not use this file except in compliance with | |
7 | * 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, software | |
12 | * distributed under the License is distributed on an "AS IS" BASIS, | |
13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
14 | * See the License for the specific language governing permissions and | |
15 | * limitations under the License. | |
16 | */ | |
17 | ||
18 | package org.apache.arrow.vector.ipc.message; | |
19 | ||
20 | import static org.apache.arrow.vector.ipc.message.FBSerializables.writeAllStructsToVector; | |
21 | import static org.apache.arrow.vector.ipc.message.FBSerializables.writeKeyValues; | |
22 | ||
23 | import java.util.ArrayList; | |
24 | import java.util.HashMap; | |
25 | import java.util.List; | |
26 | import java.util.Map; | |
27 | ||
28 | import org.apache.arrow.flatbuf.Block; | |
29 | import org.apache.arrow.flatbuf.Footer; | |
30 | import org.apache.arrow.flatbuf.KeyValue; | |
31 | import org.apache.arrow.vector.types.MetadataVersion; | |
32 | import org.apache.arrow.vector.types.pojo.Schema; | |
33 | ||
34 | import com.google.flatbuffers.FlatBufferBuilder; | |
35 | ||
36 | /** Footer metadata for the arrow file format. */ | |
37 | public class ArrowFooter implements FBSerializable { | |
38 | ||
39 | private final Schema schema; | |
40 | ||
41 | private final List<ArrowBlock> dictionaries; | |
42 | ||
43 | private final List<ArrowBlock> recordBatches; | |
44 | ||
45 | private final Map<String, String> metaData; | |
46 | ||
47 | private final MetadataVersion metadataVersion; | |
48 | ||
49 | public ArrowFooter(Schema schema, List<ArrowBlock> dictionaries, List<ArrowBlock> recordBatches) { | |
50 | this(schema, dictionaries, recordBatches, null); | |
51 | } | |
52 | ||
53 | /** | |
54 | * Constructs a new instance. | |
55 | * | |
56 | * @param schema The schema for record batches in the file. | |
57 | * @param dictionaries The dictionaries relevant to the file. | |
58 | * @param recordBatches The recordBatches written to the file. | |
59 | * @param metaData user-defined k-v meta data. | |
60 | */ | |
61 | public ArrowFooter( | |
62 | Schema schema, | |
63 | List<ArrowBlock> dictionaries, | |
64 | List<ArrowBlock> recordBatches, | |
65 | Map<String, String> metaData) { | |
66 | this(schema, dictionaries, recordBatches, metaData, MetadataVersion.DEFAULT); | |
67 | } | |
68 | ||
69 | /** | |
70 | * Constructs a new instance. | |
71 | * | |
72 | * @param schema The schema for record batches in the file. | |
73 | * @param dictionaries The dictionaries relevant to the file. | |
74 | * @param recordBatches The recordBatches written to the file. | |
75 | * @param metaData user-defined k-v meta data. | |
76 | * @param metadataVersion The Arrow metadata version. | |
77 | */ | |
78 | public ArrowFooter( | |
79 | Schema schema, | |
80 | List<ArrowBlock> dictionaries, | |
81 | List<ArrowBlock> recordBatches, | |
82 | Map<String, String> metaData, | |
83 | MetadataVersion metadataVersion) { | |
84 | this.schema = schema; | |
85 | this.dictionaries = dictionaries; | |
86 | this.recordBatches = recordBatches; | |
87 | this.metaData = metaData; | |
88 | this.metadataVersion = metadataVersion; | |
89 | } | |
90 | ||
91 | /** | |
92 | * Constructs from the corresponding Flatbuffer message. | |
93 | */ | |
94 | public ArrowFooter(Footer footer) { | |
95 | this( | |
96 | Schema.convertSchema(footer.schema()), | |
97 | dictionaries(footer), | |
98 | recordBatches(footer), | |
99 | metaData(footer), | |
100 | MetadataVersion.fromFlatbufID(footer.version()) | |
101 | ); | |
102 | } | |
103 | ||
104 | private static List<ArrowBlock> recordBatches(Footer footer) { | |
105 | List<ArrowBlock> recordBatches = new ArrayList<>(); | |
106 | Block tempBlock = new Block(); | |
107 | int recordBatchesLength = footer.recordBatchesLength(); | |
108 | for (int i = 0; i < recordBatchesLength; i++) { | |
109 | Block block = footer.recordBatches(tempBlock, i); | |
110 | recordBatches.add(new ArrowBlock(block.offset(), block.metaDataLength(), block.bodyLength())); | |
111 | } | |
112 | return recordBatches; | |
113 | } | |
114 | ||
115 | private static List<ArrowBlock> dictionaries(Footer footer) { | |
116 | List<ArrowBlock> dictionaries = new ArrayList<>(); | |
117 | Block tempBlock = new Block(); | |
118 | ||
119 | int dictionariesLength = footer.dictionariesLength(); | |
120 | for (int i = 0; i < dictionariesLength; i++) { | |
121 | Block block = footer.dictionaries(tempBlock, i); | |
122 | dictionaries.add(new ArrowBlock(block.offset(), block.metaDataLength(), block.bodyLength())); | |
123 | } | |
124 | return dictionaries; | |
125 | } | |
126 | ||
127 | private static Map<String, String> metaData(Footer footer) { | |
128 | Map<String, String> metaData = new HashMap<>(); | |
129 | ||
130 | int metaDataLength = footer.customMetadataLength(); | |
131 | for (int i = 0; i < metaDataLength; i++) { | |
132 | KeyValue kv = footer.customMetadata(i); | |
133 | metaData.put(kv.key(), kv.value()); | |
134 | } | |
135 | ||
136 | return metaData; | |
137 | } | |
138 | ||
139 | public Schema getSchema() { | |
140 | return schema; | |
141 | } | |
142 | ||
143 | public List<ArrowBlock> getDictionaries() { | |
144 | return dictionaries; | |
145 | } | |
146 | ||
147 | public List<ArrowBlock> getRecordBatches() { | |
148 | return recordBatches; | |
149 | } | |
150 | ||
151 | public Map<String, String> getMetaData() { | |
152 | return metaData; | |
153 | } | |
154 | ||
155 | public MetadataVersion getMetadataVersion() { | |
156 | return metadataVersion; | |
157 | } | |
158 | ||
159 | @Override | |
160 | public int writeTo(FlatBufferBuilder builder) { | |
161 | int schemaIndex = schema.getSchema(builder); | |
162 | Footer.startDictionariesVector(builder, dictionaries.size()); | |
163 | int dicsOffset = writeAllStructsToVector(builder, dictionaries); | |
164 | Footer.startRecordBatchesVector(builder, recordBatches.size()); | |
165 | int rbsOffset = writeAllStructsToVector(builder, recordBatches); | |
166 | ||
167 | int metaDataOffset = 0; | |
168 | if (metaData != null) { | |
169 | metaDataOffset = writeKeyValues(builder, metaData); | |
170 | } | |
171 | ||
172 | Footer.startFooter(builder); | |
173 | Footer.addSchema(builder, schemaIndex); | |
174 | Footer.addDictionaries(builder, dicsOffset); | |
175 | Footer.addRecordBatches(builder, rbsOffset); | |
176 | Footer.addCustomMetadata(builder, metaDataOffset); | |
177 | Footer.addVersion(builder, metadataVersion.toFlatbufID()); | |
178 | return Footer.endFooter(builder); | |
179 | } | |
180 | ||
181 | @Override | |
182 | public int hashCode() { | |
183 | final int prime = 31; | |
184 | int result = 1; | |
185 | result = prime * result + ((dictionaries == null) ? 0 : dictionaries.hashCode()); | |
186 | result = prime * result + ((recordBatches == null) ? 0 : recordBatches.hashCode()); | |
187 | result = prime * result + ((schema == null) ? 0 : schema.hashCode()); | |
188 | return result; | |
189 | } | |
190 | ||
191 | @Override | |
192 | public boolean equals(Object obj) { | |
193 | if (this == obj) { | |
194 | return true; | |
195 | } | |
196 | if (obj == null) { | |
197 | return false; | |
198 | } | |
199 | if (getClass() != obj.getClass()) { | |
200 | return false; | |
201 | } | |
202 | ArrowFooter other = (ArrowFooter) obj; | |
203 | if (dictionaries == null) { | |
204 | if (other.dictionaries != null) { | |
205 | return false; | |
206 | } | |
207 | } else if (!dictionaries.equals(other.dictionaries)) { | |
208 | return false; | |
209 | } | |
210 | if (recordBatches == null) { | |
211 | if (other.recordBatches != null) { | |
212 | return false; | |
213 | } | |
214 | } else if (!recordBatches.equals(other.recordBatches)) { | |
215 | return false; | |
216 | } | |
217 | if (schema == null) { | |
218 | if (other.schema != null) { | |
219 | return false; | |
220 | } | |
221 | } else if (!schema.equals(other.schema)) { | |
222 | return false; | |
223 | } | |
224 | return true; | |
225 | } | |
226 | } |