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 import org.apache.arrow.memory.ArrowBuf; 19 import org.apache.arrow.vector.complex.writer.Decimal256Writer; 20 import org.apache.arrow.vector.complex.writer.DecimalWriter; 21 import org.apache.arrow.vector.holders.Decimal256Holder; 22 import org.apache.arrow.vector.holders.DecimalHolder; 23 24 import java.lang.UnsupportedOperationException; 25 import java.math.BigDecimal; 26 27 <@pp.dropOutputFile /> 28 <@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionMapWriter.java" /> 29 30 31 <#include "/@includes/license.ftl" /> 32 33 package org.apache.arrow.vector.complex.impl; 34 35 <#include "/@includes/vv_imports.ftl" /> 36 37 /* 38 * This class is generated using freemarker and the ${.template_name} template. 39 */ 40 41 /** 42 * <p>Writer for MapVectors. This extends UnionListWriter to simplify writing map entries to a list 43 * of struct elements, with "key" and "value" fields. The procedure for writing a map begin with 44 * {@link #startMap()} followed by {@link #startEntry()}. An entry is written by using the 45 * {@link #key()} writer to write the key, then the {@link #value()} writer to write a value. After 46 * writing the value, call {@link #endEntry()} to complete the entry. Each map can have 1 or more 47 * entries. When done writing entries, call {@link #endMap()} to complete the map. 48 * 49 * <p>NOTE: the MapVector can have NULL values by not writing to position. If a map is started with 50 * {@link #startMap()}, then it must have a key written. The value of a map entry can be NULL by 51 * not using the {@link #value()} writer. 52 * 53 * <p>Example to write the following map to position 5 of a vector 54 * <pre>{@code 55 * // { 56 * // 1 -> 3, 57 * // 2 -> 4, 58 * // 3 -> NULL 59 * // } 60 * 61 * UnionMapWriter writer = ... 62 * 63 * writer.setPosition(5); 64 * writer.startMap(); 65 * writer.startEntry(); 66 * writer.key().integer().writeInt(1); 67 * writer.value().integer().writeInt(3); 68 * writer.endEntry(); 69 * writer.startEntry(); 70 * writer.key().integer().writeInt(2); 71 * writer.value().integer().writeInt(4); 72 * writer.endEntry(); 73 * writer.startEntry(); 74 * writer.key().integer().writeInt(3); 75 * writer.endEntry(); 76 * writer.endMap(); 77 * </pre> 78 * </p> 79 */ 80 @SuppressWarnings("unused") 81 public class UnionMapWriter extends UnionListWriter { 82 83 /** 84 * Current mode for writing map entries, set by calling {@link #key()} or {@link #value()} 85 * and reset with a call to {@link #endEntry()}. With KEY mode, a struct writer with field 86 * named "key" is returned. With VALUE mode, a struct writer with field named "value" is 87 * returned. In OFF mode, the writer will behave like a standard UnionListWriter 88 */ 89 private enum MapWriteMode { 90 OFF, 91 KEY, 92 VALUE, 93 } 94 95 private MapWriteMode mode = MapWriteMode.OFF; 96 private StructWriter entryWriter; 97 UnionMapWriter(MapVector vector)98 public UnionMapWriter(MapVector vector) { 99 super(vector); 100 entryWriter = struct(); 101 } 102 103 /** Start writing a map that consists of 1 or more entries. */ startMap()104 public void startMap() { 105 startList(); 106 } 107 108 /** Complete the map. */ endMap()109 public void endMap() { 110 endList(); 111 } 112 113 /** 114 * Start a map entry that should be followed by calls to {@link #key()} and {@link #value()} 115 * writers. Call {@link #endEntry()} to complete the entry. 116 */ startEntry()117 public void startEntry() { 118 writer.setAddVectorAsNullable(false); 119 entryWriter.start(); 120 } 121 122 /** Complete the map entry. */ endEntry()123 public void endEntry() { 124 entryWriter.end(); 125 mode = MapWriteMode.OFF; 126 writer.setAddVectorAsNullable(true); 127 } 128 129 /** Return the key writer that is used to write to the "key" field. */ key()130 public UnionMapWriter key() { 131 writer.setAddVectorAsNullable(false); 132 mode = MapWriteMode.KEY; 133 return this; 134 } 135 136 /** Return the value writer that is used to write to the "value" field. */ value()137 public UnionMapWriter value() { 138 writer.setAddVectorAsNullable(true); 139 mode = MapWriteMode.VALUE; 140 return this; 141 } 142 143 <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> 144 <#assign fields = minor.fields!type.fields /> 145 <#assign uncappedName = name?uncap_first/> 146 <#if uncappedName == "int" ><#assign uncappedName = "integer" /></#if> 147 <#if !minor.typeParams?? > 148 @Override 149 public ${name}Writer ${uncappedName}() { 150 switch (mode) { 151 case KEY: 152 return entryWriter.${uncappedName}(MapVector.KEY_NAME); 153 case VALUE: 154 return entryWriter.${uncappedName}(MapVector.VALUE_NAME); 155 default: 156 return this; 157 } 158 } 159 160 </#if> 161 </#list></#list> 162 @Override decimal()163 public DecimalWriter decimal() { 164 switch (mode) { 165 case KEY: 166 return entryWriter.decimal(MapVector.KEY_NAME); 167 case VALUE: 168 return entryWriter.decimal(MapVector.VALUE_NAME); 169 default: 170 return this; 171 } 172 } 173 174 @Override decimal256()175 public Decimal256Writer decimal256() { 176 switch (mode) { 177 case KEY: 178 return entryWriter.decimal256(MapVector.KEY_NAME); 179 case VALUE: 180 return entryWriter.decimal256(MapVector.VALUE_NAME); 181 default: 182 return this; 183 } 184 } 185 186 187 @Override struct()188 public StructWriter struct() { 189 switch (mode) { 190 case KEY: 191 return entryWriter.struct(MapVector.KEY_NAME); 192 case VALUE: 193 return entryWriter.struct(MapVector.VALUE_NAME); 194 default: 195 return super.struct(); 196 } 197 } 198 199 @Override list()200 public ListWriter list() { 201 switch (mode) { 202 case KEY: 203 return entryWriter.list(MapVector.KEY_NAME); 204 case VALUE: 205 return entryWriter.list(MapVector.VALUE_NAME); 206 default: 207 return super.list(); 208 } 209 } 210 211 @Override map(boolean keysSorted)212 public MapWriter map(boolean keysSorted) { 213 switch (mode) { 214 case KEY: 215 return entryWriter.map(MapVector.KEY_NAME, keysSorted); 216 case VALUE: 217 return entryWriter.map(MapVector.VALUE_NAME, keysSorted); 218 default: 219 return super.map(); 220 } 221 } 222 } 223