support root-level maps

Author: cowtowncoder

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/AvroWriteContext.java

@@ -146,10 +146,12 @@ protected GenericArray<Object> _createArray(Schema schema)
 
     protected AvroWriteContext _createObjectContext(Schema schema) throws JsonMappingException
     {
-        if (schema.getType() == Schema.Type.UNION) {
+        Type type = schema.getType();
+        if (type == Schema.Type.UNION) {
             schema = _recordOrMapFromUnion(schema);
+            type = schema.getType();
         }
-        if (schema.getType() == Schema.Type.MAP) {
+        if (type == Schema.Type.MAP) {
             return new MapWriteContext(this, _generator, schema);
         }
         return new ObjectWriteContext(this, _generator, _createRecord(schema));

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/MapWriteContext.java

@@ -16,7 +16,7 @@ public final class MapWriteContext
     extends KeyValueContext
 {
     protected final Map<String,Object> _data;
-    
+
     public MapWriteContext(AvroWriteContext parent, AvroGenerator generator, Schema schema)
     {
         super(parent, generator, schema);

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/RootContext.java

@@ -13,13 +13,17 @@ class RootContext
     extends AvroWriteContext
 {
     /**
-     * We need to keep reference to the root value here.
-     *<p>
-     * TODO: What about Map values at root? Avro codec does not seem
-     * able to support it via Generic API.
+     * We need to keep reference to the root value here; either
+     * <code>GenericContainer</code> or <code>Map</code> (yes,
+     * Avro APIs are... odd).
      */
-    protected GenericContainer _rootValue;
-    
+    protected Object _rootValue;
+
+    /**
+     * Lazily created instance for encoding: reused in case of root value sequences.
+     */
+    private NonBSGenericDatumWriter<Object> _writer;
+
     public RootContext(AvroGenerator generator, Schema schema) {
         super(TYPE_ROOT, null, generator, schema);
     }
@@ -51,16 +55,21 @@ public final AvroWriteContext createChildObjectContext() throws JsonMappingExcep
         switch (_schema.getType()) {
         case RECORD:
         case UNION: // maybe
-            break;
-        case MAP:
-            throw new UnsupportedOperationException("Root-level Maps not supported: Avro Codec has no way to create these");
+            {
+                GenericRecord rec = _createRecord(_schema);
+                _rootValue = rec;
+                return new ObjectWriteContext(this, _generator, rec);
+            }
+        case MAP: // used to not be supported
+            {
+                MapWriteContext ctxt = new MapWriteContext(this, _generator, _schema);
+                _rootValue = ctxt.rawValue();
+                return ctxt;
+            }
         default:
-            throw new IllegalStateException("Can not write START_OBJECT; schema type is "
-                    +_schema.getType());
         }
-        GenericRecord rec = _createRecord(_schema);
-        _rootValue = rec;
-        return new ObjectWriteContext(this, _generator, rec);
+        throw new IllegalStateException("Can not write START_OBJECT; schema type is "
+                +_schema.getType());
     }
 
     @Override
@@ -73,15 +82,10 @@ public void writeString(String value) {
         _reportError();
     }
 
-    /**
-     * Lazily created instance for encoding: reused in case of root value sequences.
-     */
-    private NonBSGenericDatumWriter<GenericContainer> _writer;
-
     @Override
     public void complete(BinaryEncoder encoder) throws IOException {
         if (_writer == null) {
-            _writer = new NonBSGenericDatumWriter<GenericContainer>(_schema);
+            _writer = new NonBSGenericDatumWriter<Object>(_schema);
         }
         _writer.write(_rootValue, encoder);
     }

avro/src/test/java/com/fasterxml/jackson/dataformat/avro/MapTest.java

@@ -1,11 +1,16 @@
 package com.fasterxml.jackson.dataformat.avro;
 
+import static org.junit.Assert.assertArrayEquals;
+
 import java.io.ByteArrayOutputStream;
 import java.util.*;
 
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
+import com.fasterxml.jackson.databind.MappingIterator;
+import com.fasterxml.jackson.databind.SequenceWriter;
+import com.fasterxml.jackson.dataformat.avro.AvroTestBase.Employee;
 
 public class MapTest extends AvroTestBase
 {
@@ -36,6 +41,8 @@ public void setStuff(Map<String,String> arg) {
         }
     }
 
+    private final AvroMapper MAPPER = getMapper();
+
     public void testRecordWithMap() throws Exception
     {
         AvroMapper mapper = getMapper();
@@ -145,16 +152,61 @@ public void testMapOrNull() throws Exception
     //   since Records and Arrays work, but looks like there are some issues
     //   regarding them so can't yet test
 
-    /*
     public void testRootStringMap() throws Exception
     {
-        AvroMapper mapper = getMapper();
         AvroSchema schema = getStringMapSchema();
-        Map<String,String> input = new LinkedHashMap<>();
-        input.put("a", "1");
-        input.put("b", "2");
+        Map<String,String> input = _map("a", "1", "b", "2");
+
+        byte[] b = MAPPER.writer(schema).writeValueAsBytes(input);
+        Map<String,String> result = MAPPER.readerFor(Map.class)
+                .with(schema)
+                .readValue(b);
+        assertEquals(2, result.size());
+        assertEquals("1", result.get("a"));
+        assertEquals("2", result.get("b"));
+    }
+
+    public void testRootMapSequence() throws Exception
+    {
+        ByteArrayOutputStream b = new ByteArrayOutputStream(1000);
+        AvroSchema schema = getStringMapSchema();
+        Map<String,String> input1 = _map("a", "1", "b", "2");
+        Map<String,String> input2 = _map("c", "3", "d", "4");
+
+        SequenceWriter sw = MAPPER.writerFor(Map.class)
+            .with(schema)
+            .writeValues(b);
+        sw.write(input1);
+        int curr = b.size();
+        sw.write(input2);
+        int diff = b.size() - curr;
+        if (diff == 0) {
+            fail("Should have output more bytes for second entry, did not, total: "+curr);
+        }
+        sw.close();
 
-        byte[] b = mapper.writer(schema).writeValueAsBytes(input);
+        byte[] bytes = b.toByteArray();
+
+        assertNotNull(bytes);
+
+        MappingIterator<Map<String,String>> it = MAPPER.readerFor(Map.class)
+                .with(schema)
+                .readValues(bytes);
+        assertTrue(it.hasNextValue());
+        assertEquals(input1, it.nextValue());
+
+        assertTrue(it.hasNextValue());
+        assertEquals(input2, it.nextValue());
+
+        assertFalse(it.hasNextValue());
+        it.close();
+    }
+
+    private Map<String,String> _map(String... stuff) {
+        Map<String,String> map = new LinkedHashMap<String,String>();
+        for (int i = 0, end = stuff.length; i < end; i += 2) {
+            map.put(stuff[i], stuff[i+1]);
+        }
+        return map;
     }
-    */
 }