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