Enable non-nested dictionary batches in Flight integration tests

cpp/src/arrow/flight/client.cc

@@ -308,7 +308,8 @@ class DoPutPayloadWriter : public ipc::internal::IpcPayloadWriter {
       }
       RETURN_NOT_OK(Buffer::FromString(str_descr, &payload.descriptor));
       first_payload_ = false;
-    } else if (stream_writer_->app_metadata_) {
+    } else if (ipc_payload.type == ipc::Message::RECORD_BATCH &&
+               stream_writer_->app_metadata_) {
       payload.app_metadata = std::move(stream_writer_->app_metadata_);
     }
 

cpp/src/arrow/flight/test-integration-server.cc

@@ -106,7 +106,9 @@ class FlightIntegrationTestServer : public FlightServerBase {
       RETURN_NOT_OK(reader->ReadWithMetadata(&chunk, &metadata));
       if (chunk == nullptr) break;
       retrieved_chunks.push_back(chunk);
-      RETURN_NOT_OK(writer->WriteMetadata(*metadata));
+      if (metadata) {
+        RETURN_NOT_OK(writer->WriteMetadata(*metadata));
+      }
     }
     std::shared_ptr<arrow::Table> retrieved_data;
     RETURN_NOT_OK(arrow::Table::FromRecordBatches(reader->schema(), retrieved_chunks,

cpp/src/arrow/flight/test-util.cc

@@ -271,8 +271,10 @@ Status NumberingStream::GetSchemaPayload(FlightPayload* payload) {
 
 Status NumberingStream::Next(FlightPayload* payload) {
   RETURN_NOT_OK(stream_->Next(payload));
-  payload->app_metadata = Buffer::FromString(std::to_string(counter_));
-  counter_++;
+  if (payload && payload->ipc_message.type == ipc::Message::RECORD_BATCH) {
+    payload->app_metadata = Buffer::FromString(std::to_string(counter_));
+    counter_++;
+  }
   return Status::OK();
 }
 

integration/integration_test.py

@@ -1152,7 +1152,7 @@ def get_generated_json_files(tempdir=None, flight=False):
         generate_interval_case(),
         generate_map_case(),
         generate_nested_case(),
-        generate_dictionary_case().skip_category(SKIP_FLIGHT),
+        generate_dictionary_case(),
         generate_nested_dictionary_case().skip_category(SKIP_ARROW)
                                          .skip_category(SKIP_FLIGHT),
     ]

java/flight/src/main/java/org/apache/arrow/flight/ArrowMessage.java

@@ -35,6 +35,7 @@ import org.apache.arrow.flight.impl.Flight.FlightData;
 import org.apache.arrow.flight.impl.Flight.FlightDescriptor;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.ipc.message.MessageSerializer;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -147,6 +148,21 @@ class ArrowMessage implements AutoCloseable {
     this.appMetadata = appMetadata;
   }
 
+  public ArrowMessage(ArrowDictionaryBatch batch) {
+    FlatBufferBuilder builder = new FlatBufferBuilder();
+    int batchOffset = batch.writeTo(builder);
+    ByteBuffer serializedMessage = MessageSerializer
+        .serializeMessage(builder, MessageHeader.DictionaryBatch, batchOffset,
+            batch.computeBodyLength());
+    serializedMessage = serializedMessage.slice();
+    this.message = Message.getRootAsMessage(serializedMessage);
+    // asInputStream will free the buffers implicitly, so increment the reference count
+    batch.getDictionary().getBuffers().forEach(buf -> buf.getReferenceManager().retain());
+    this.bufs = ImmutableList.copyOf(batch.getDictionary().getBuffers());
+    this.descriptor = null;
+    this.appMetadata = null;
+  }
+
   private ArrowMessage(FlightDescriptor descriptor, Message message, byte[] appMetadata, ArrowBuf buf) {
     this.message = message;
     this.descriptor = descriptor;
@@ -184,6 +200,13 @@ class ArrowMessage implements AutoCloseable {
     return batch;
   }
 
+  public ArrowDictionaryBatch asDictionaryBatch() throws IOException {
+    Preconditions.checkArgument(bufs.size() == 1, "A batch can only be consumed if it contains a single ArrowBuf.");
+    Preconditions.checkArgument(getMessageType() == HeaderType.DICTIONARY_BATCH);
+    ArrowBuf underlying = bufs.get(0);
+    return MessageSerializer.deserializeDictionaryBatch(message, underlying);
+  }
+
   public Iterable<ArrowBuf> getBufs() {
     return Iterables.unmodifiableIterable(bufs);
   }
@@ -277,7 +300,8 @@ class ArrowMessage implements AutoCloseable {
         return NO_BODY_MARSHALLER.stream(builder.build());
       }
 
-      Preconditions.checkArgument(getMessageType() == HeaderType.RECORD_BATCH);
+      Preconditions.checkArgument(getMessageType() == HeaderType.RECORD_BATCH ||
+          getMessageType() == HeaderType.DICTIONARY_BATCH);
       Preconditions.checkArgument(!bufs.isEmpty());
       Preconditions.checkArgument(descriptor == null, "Descriptor should only be included in the schema message.");
 

java/flight/src/main/java/org/apache/arrow/flight/DictionaryUtils.java

@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.function.Consumer;
+import java.util.stream.Stream;
+
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.DictionaryUtility;
+
+/**
+ * Utilities to work with dictionaries in Flight.
+ */
+final class DictionaryUtils {
+
+  private DictionaryUtils() {
+    throw new UnsupportedOperationException("Do not instantiate this class.");
+  }
+
+  /**
+   * Generate all the necessary Flight messages to send a schema and associated dictionaries.
+   */
+  static Schema generateSchemaMessages(final Schema originalSchema, final FlightDescriptor descriptor,
+      final DictionaryProvider provider, final Consumer<ArrowMessage> messageCallback) {
+    final List<Field> fields = new ArrayList<>(originalSchema.getFields().size());
+    final Set<Long> dictionaryIds = new HashSet<>();
+    for (final Field field : originalSchema.getFields()) {
+      fields.add(DictionaryUtility.toMessageFormat(field, provider, dictionaryIds));
+    }
+    final Schema schema = new Schema(fields, originalSchema.getCustomMetadata());
+    // Send the schema message
+    messageCallback.accept(new ArrowMessage(descriptor == null ? null : descriptor.toProtocol(), schema));
+    // Create and write dictionary batches
+    for (Long id : dictionaryIds) {
+      final Dictionary dictionary = provider.lookup(id);
+      final FieldVector vector = dictionary.getVector();
+      final int count = vector.getValueCount();
+      // Do NOT close this root, as it does not actually own the vector.
+      final VectorSchemaRoot dictRoot = new VectorSchemaRoot(
+          Collections.singletonList(vector.getField()),
+          Collections.singletonList(vector),
+          count);
+      final VectorUnloader unloader = new VectorUnloader(dictRoot);
+      try (final ArrowDictionaryBatch dictionaryBatch = new ArrowDictionaryBatch(
+          id, unloader.getRecordBatch())) {
+        messageCallback.accept(new ArrowMessage(dictionaryBatch));
+      }
+    }
+    return schema;
+  }
+}

java/flight/src/main/java/org/apache/arrow/flight/FlightClient.java

@@ -19,7 +19,12 @@ package org.apache.arrow.flight;
 
 import java.io.InputStream;
 import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -37,9 +42,17 @@ import org.apache.arrow.flight.impl.FlightServiceGrpc;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceBlockingStub;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceStub;
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider;
+import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.DictionaryUtility;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
@@ -164,6 +177,19 @@ public class FlightClient implements AutoCloseable {
    */
   public ClientStreamListener startPut(FlightDescriptor descriptor, VectorSchemaRoot root,
       StreamListener<PutResult> metadataListener, CallOption... options) {
+    return startPut(descriptor, root, new MapDictionaryProvider(), metadataListener, options);
+  }
+
+  /**
+   * Create or append a descriptor with another stream.
+   * @param descriptor FlightDescriptor the descriptor for the data
+   * @param root VectorSchemaRoot the root containing data
+   * @param metadataListener A handler for metadata messages from the server.
+   * @param options RPC-layer hints for this call.
+   * @return ClientStreamListener an interface to control uploading data
+   */
+  public ClientStreamListener startPut(FlightDescriptor descriptor, VectorSchemaRoot root, DictionaryProvider provider,
+      StreamListener<PutResult> metadataListener, CallOption... options) {
     Preconditions.checkNotNull(descriptor);
     Preconditions.checkNotNull(root);
 
@@ -173,8 +199,7 @@ public class FlightClient implements AutoCloseable {
         ClientCalls.asyncBidiStreamingCall(
             authInterceptor.interceptCall(doPutDescriptor, callOptions, channel), resultObserver);
     // send the schema to start.
-    ArrowMessage message = new ArrowMessage(descriptor.toProtocol(), root.getSchema());
-    observer.onNext(message);
+    DictionaryUtils.generateSchemaMessages(root.getSchema(), descriptor, provider, observer::onNext);
     return new PutObserver(new VectorUnloader(
         root, true /* include # of nulls in vectors */, true /* must align buffers to be C++-compatible */),
         observer, resultObserver.getFuture());

java/flight/src/main/java/org/apache/arrow/flight/FlightProducer.java

@@ -18,6 +18,7 @@
 package org.apache.arrow.flight;
 
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
 
 /**
  * API to Implement an Arrow Flight producer.
@@ -100,6 +101,13 @@ public interface FlightProducer {
      */
     void start(VectorSchemaRoot root);
 
+    /**
+     * Start sending data, using the schema of the given {@link VectorSchemaRoot}.
+     *
+     * <p>This method must be called before all others.
+     */
+    void start(VectorSchemaRoot root, DictionaryProvider dictionaries);
+
     /**
      * Send the current contents of the associated {@link VectorSchemaRoot}.
      */

java/flight/src/main/java/org/apache/arrow/flight/FlightService.java

@@ -17,6 +17,11 @@
 
 package org.apache.arrow.flight;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.function.BooleanSupplier;
@@ -32,8 +37,17 @@ import org.apache.arrow.flight.impl.Flight.HandshakeRequest;
 import org.apache.arrow.flight.impl.Flight.HandshakeResponse;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider;
+import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.DictionaryUtility;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -137,9 +151,14 @@ class FlightService extends FlightServiceImplBase {
 
     @Override
     public void start(VectorSchemaRoot root) {
-      responseObserver.onNext(new ArrowMessage(null, root.getSchema()));
-      // [ARROW-4213] We must align buffers to be compatible with other languages.
+      start(root, new MapDictionaryProvider());
+    }
+
+    @Override
+    public void start(VectorSchemaRoot root, DictionaryProvider provider) {
       unloader = new VectorUnloader(root, true, true);
+
+      DictionaryUtils.generateSchemaMessages(root.getSchema(), null, provider, responseObserver::onNext);
     }
 
     @Override
@@ -171,7 +190,7 @@ class FlightService extends FlightServiceImplBase {
     responseObserver.disableAutoInboundFlowControl();
     responseObserver.request(1);
 
-    FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, null, (count) -> responseObserver.request(count));
+    FlightStream fs = new FlightStream(allocator, PENDING_REQUESTS, null, responseObserver::request);
     executors.submit(() -> {
       try {
         producer.acceptPut(makeContext(responseObserver), fs,
@@ -179,6 +198,9 @@ class FlightService extends FlightServiceImplBase {
         responseObserver.onCompleted();
       } catch (Exception ex) {
         responseObserver.onError(ex);
+        // The client may have terminated, so the exception here is effectively swallowed.
+        // Log the error as well so -something- makes it to the developer.
+        logger.error("Exception handling DoPut", ex);
       }
     });
 

java/flight/src/main/java/org/apache/arrow/flight/FlightStream.java

@@ -17,17 +17,28 @@
 
 package org.apache.arrow.flight;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.stream.Collectors;
 
+import org.apache.arrow.flight.ArrowMessage.HeaderType;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorLoader;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.arrow.vector.util.DictionaryUtility;
 
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
@@ -56,6 +67,7 @@ public class FlightStream implements AutoCloseable {
   private volatile int pending = 1;
   private boolean completed = false;
   private volatile VectorSchemaRoot fulfilledRoot;
+  private DictionaryProvider.MapDictionaryProvider dictionaries;
   private volatile VectorLoader loader;
   private volatile Throwable ex;
   private volatile FlightDescriptor descriptor;
@@ -75,12 +87,17 @@ public class FlightStream implements AutoCloseable {
     this.pendingTarget = pendingTarget;
     this.cancellable = cancellable;
     this.requestor = requestor;
+    this.dictionaries = new DictionaryProvider.MapDictionaryProvider();
   }
 
   public Schema getSchema() {
     return schema;
   }
 
+  public DictionaryProvider getDictionaryProvider() {
+    return dictionaries;
+  }
+
   public FlightDescriptor getDescriptor() {
     return descriptor;
   }
@@ -133,15 +150,33 @@ public class FlightStream implements AutoCloseable {
         }
       } else {
         ArrowMessage msg = ((ArrowMessage) data);
-        try (ArrowRecordBatch arb = msg.asRecordBatch()) {
-          loader.load(arb);
+        if (msg.getMessageType() == HeaderType.RECORD_BATCH) {
+          try (ArrowRecordBatch arb = msg.asRecordBatch()) {
+            loader.load(arb);
+          }
+          this.applicationMetadata = msg.getApplicationMetadata();
+        } else if (msg.getMessageType() == HeaderType.DICTIONARY_BATCH) {
+          try (ArrowDictionaryBatch arb = msg.asDictionaryBatch()) {
+            final long id = arb.getDictionaryId();
+            final Dictionary dictionary = dictionaries.lookup(id);
+            if (dictionary == null) {
+              throw new IllegalArgumentException("Dictionary not defined in schema: ID " + id);
+            }
+
+            final FieldVector vector = dictionary.getVector();
+            final VectorSchemaRoot dictionaryRoot = new VectorSchemaRoot(Collections.singletonList(vector.getField()),
+                Collections.singletonList(vector), 0);
+            final VectorLoader dictionaryLoader = new VectorLoader(dictionaryRoot);
+            dictionaryLoader.load(arb.getDictionary());
+          }
+          return next();
+        } else {
+          throw new UnsupportedOperationException("Message type is unsupported: " + msg.getMessageType());
         }
-        this.applicationMetadata = msg.getApplicationMetadata();
         return true;
       }
-
     } catch (Exception e) {
-      throw Throwables.propagate(e);
+      throw new RuntimeException(e);
     }
   }
 
@@ -181,23 +216,36 @@ public class FlightStream implements AutoCloseable {
     public void onNext(ArrowMessage msg) {
       requestOutstanding();
       switch (msg.getMessageType()) {
-        case SCHEMA:
+        case SCHEMA: {
           schema = msg.asSchema();
+          final List<Field> fields = new ArrayList<>();
+          final Map<Long, Dictionary> dictionaryMap = new HashMap<>();
+          for (final Field originalField : schema.getFields()) {
+            final Field updatedField = DictionaryUtility.toMemoryFormat(originalField, allocator, dictionaryMap);
+            fields.add(updatedField);
+          }
+          for (final Map.Entry<Long, Dictionary> entry : dictionaryMap.entrySet()) {
+            dictionaries.put(entry.getValue());
+          }
+          schema = new Schema(fields, schema.getCustomMetadata());
           fulfilledRoot = VectorSchemaRoot.create(schema, allocator);
           loader = new VectorLoader(fulfilledRoot);
           descriptor = msg.getDescriptor() != null ? new FlightDescriptor(msg.getDescriptor()) : null;
           root.set(fulfilledRoot);
 
           break;
+        }
         case RECORD_BATCH:
           queue.add(msg);
           break;
-        case NONE:
         case DICTIONARY_BATCH:
+          queue.add(msg);
+          break;
+        case NONE:
         case TENSOR:
         default:
           queue.add(DONE_EX);
-          ex = new UnsupportedOperationException("Unable to handle message of type: " + msg);
+          ex = new UnsupportedOperationException("Unable to handle message of type: " + msg.getMessageType());
 
       }
 

java/flight/src/main/java/org/apache/arrow/flight/GenericOperation.java

@@ -1,42 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.arrow.flight;
-
-/**
- * Unused?.
- */
-class GenericOperation {
-
-  private final String type;
-  private final byte[] body;
-
-  public GenericOperation(String type, byte[] body) {
-    super();
-    this.type = type;
-    this.body = body == null ? new byte[0] : body;
-  }
-
-  public String getType() {
-    return type;
-  }
-
-  public byte[] getBody() {
-    return body;
-  }
-
-}

java/flight/src/main/java/org/apache/arrow/flight/PutResult.java

@@ -36,6 +36,9 @@ public class PutResult {
 
   /** Create a PutResult with application-specific metadata. */
   public static PutResult metadata(byte[] metadata) {
+    if (metadata == null) {
+      return empty();
+    }
     return new PutResult(ByteBuffer.wrap(metadata));
   }
 

java/flight/src/main/java/org/apache/arrow/flight/example/FlightHolder.java

@@ -28,6 +28,7 @@ import org.apache.arrow.flight.FlightInfo;
 import org.apache.arrow.flight.Location;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.util.AutoCloseables;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.types.pojo.Schema;
 
 import com.google.common.base.Preconditions;
@@ -43,19 +44,22 @@ public class FlightHolder implements AutoCloseable {
   private final FlightDescriptor descriptor;
   private final Schema schema;
   private final List<Stream> streams = new CopyOnWriteArrayList<>();
+  private final DictionaryProvider dictionaryProvider;
 
   /**
    * Creates a new instance.
-   *
-   * @param allocator The allocator to use for allocating buffers to store data.
+   *  @param allocator The allocator to use for allocating buffers to store data.
    * @param descriptor The descriptor for the streams.
    * @param schema  The schema for the stream.
+   * @param dictionaryProvider The dictionary provider for the stream.
    */
-  public FlightHolder(BufferAllocator allocator, FlightDescriptor descriptor, Schema schema) {
+  public FlightHolder(BufferAllocator allocator, FlightDescriptor descriptor, Schema schema,
+      DictionaryProvider dictionaryProvider) {
     Preconditions.checkArgument(!descriptor.isCommand());
     this.allocator = allocator.newChildAllocator(descriptor.toString(), 0, Long.MAX_VALUE);
     this.descriptor = descriptor;
     this.schema = schema;
+    this.dictionaryProvider = dictionaryProvider;
   }
 
   /**
@@ -72,8 +76,8 @@ public class FlightHolder implements AutoCloseable {
    * Adds a new streams which clients can populate via the returned object.
    */
   public Stream.StreamCreator addStream(Schema schema) {
-    Preconditions.checkArgument(schema.equals(schema), "Stream schema inconsistent with existing schema.");
-    return new Stream.StreamCreator(schema, allocator, t -> {
+    Preconditions.checkArgument(this.schema.equals(schema), "Stream schema inconsistent with existing schema.");
+    return new Stream.StreamCreator(schema, dictionaryProvider, allocator, t -> {
       synchronized (streams) {
         streams.add(t);
       }

java/flight/src/main/java/org/apache/arrow/flight/example/InMemoryStore.java

@@ -79,17 +79,6 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
     return h.getStream(example);
   }
 
-  /**
-   * Create a new {@link Stream} with the given schema and descriptor.
-   */
-  public StreamCreator putStream(final FlightDescriptor descriptor, final Schema schema) {
-    final FlightHolder h = holders.computeIfAbsent(
-        descriptor,
-        t -> new FlightHolder(allocator, t, schema));
-
-    return h.addStream(schema);
-  }
-
   @Override
   public void listFlights(CallContext context, Criteria criteria,
       StreamListener<FlightInfo> listener) {
@@ -123,7 +112,7 @@ public class InMemoryStore implements FlightProducer, AutoCloseable {
       try (VectorSchemaRoot root = flightStream.getRoot()) {
         final FlightHolder h = holders.computeIfAbsent(
             flightStream.getDescriptor(),
-            t -> new FlightHolder(allocator, t, flightStream.getSchema()));
+            t -> new FlightHolder(allocator, t, flightStream.getSchema(), flightStream.getDictionaryProvider()));
 
         creator = h.addStream(flightStream.getSchema());
 

java/flight/src/main/java/org/apache/arrow/flight/example/Stream.java

@@ -29,6 +29,7 @@ import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.VectorLoader;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.types.pojo.Schema;
 
@@ -41,6 +42,7 @@ import com.google.common.collect.ImmutableList;
 public class Stream implements AutoCloseable, Iterable<ArrowRecordBatch> {
 
   private final String uuid = UUID.randomUUID().toString();
+  private final DictionaryProvider dictionaryProvider;
   private final List<ArrowRecordBatch> batches;
   private final Schema schema;
   private final long recordCount;
@@ -54,9 +56,11 @@ public class Stream implements AutoCloseable, Iterable<ArrowRecordBatch> {
    */
   public Stream(
       final Schema schema,
+      final DictionaryProvider dictionaryProvider,
       List<ArrowRecordBatch> batches,
       long recordCount) {
     this.schema = schema;
+    this.dictionaryProvider = dictionaryProvider;
     this.batches = ImmutableList.copyOf(batches);
     this.recordCount = recordCount;
   }
@@ -83,7 +87,7 @@ public class Stream implements AutoCloseable, Iterable<ArrowRecordBatch> {
    */
   public void sendTo(BufferAllocator allocator, ServerStreamListener listener) {
     try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
-      listener.start(root);
+      listener.start(root, dictionaryProvider);
       final VectorLoader loader = new VectorLoader(root);
       int counter = 0;
       for (ArrowRecordBatch batch : batches) {
@@ -121,18 +125,22 @@ public class Stream implements AutoCloseable, Iterable<ArrowRecordBatch> {
     private final List<ArrowRecordBatch> batches = new ArrayList<>();
     private final Consumer<Stream> committer;
     private long recordCount = 0;
+    private DictionaryProvider dictionaryProvider;
 
     /**
      * Creates a new instance.
      *
      * @param schema The schema for batches in the stream.
+     * @param dictionaryProvider The dictionary provider for the stream.
      * @param allocator  The allocator used to copy data permanently into the stream.
      * @param committer A callback for when the the stream is ready to be finalized (no more batches).
      */
-    public StreamCreator(Schema schema, BufferAllocator allocator, Consumer<Stream> committer) {
+    public StreamCreator(Schema schema, DictionaryProvider dictionaryProvider,
+        BufferAllocator allocator, Consumer<Stream> committer) {
       this.allocator = allocator;
       this.committer = committer;
       this.schema = schema;
+      this.dictionaryProvider = dictionaryProvider;
     }
 
     /**
@@ -155,7 +163,7 @@ public class Stream implements AutoCloseable, Iterable<ArrowRecordBatch> {
      * Complete building the stream (no more batches can be added).
      */
     public void complete() {
-      Stream stream = new Stream(schema, batches, recordCount);
+      Stream stream = new Stream(schema, dictionaryProvider, batches, recordCount);
       committer.accept(stream);
     }
 

java/flight/src/main/java/org/apache/arrow/flight/example/integration/IntegrationTestClient.java

@@ -96,27 +96,28 @@ class IntegrationTestClient {
       jsonRoot = VectorSchemaRoot.create(root.getSchema(), allocator);
       VectorUnloader unloader = new VectorUnloader(root);
       VectorLoader jsonLoader = new VectorLoader(jsonRoot);
-      FlightClient.ClientStreamListener stream = client.startPut(descriptor, root, new StreamListener<PutResult>() {
-        int counter = 0;
+      FlightClient.ClientStreamListener stream = client.startPut(descriptor, root, reader,
+          new StreamListener<PutResult>() {
+            int counter = 0;
 
-        @Override
-        public void onNext(PutResult val) {
-          final String metadata = StandardCharsets.UTF_8.decode(val.getApplicationMetadata()).toString();
-          if (!Integer.toString(counter).equals(metadata)) {
-            throw new RuntimeException(
-                String.format("Invalid ACK from server. Expected '%d' but got '%s'.", counter, metadata));
-          }
-          counter++;
-        }
+            @Override
+            public void onNext(PutResult val) {
+              final String metadata = StandardCharsets.UTF_8.decode(val.getApplicationMetadata()).toString();
+              if (!Integer.toString(counter).equals(metadata)) {
+                throw new RuntimeException(
+                    String.format("Invalid ACK from server. Expected '%d' but got '%s'.", counter, metadata));
+              }
+              counter++;
+            }
 
-        @Override
-        public void onError(Throwable t) {
-        }
+            @Override
+            public void onError(Throwable t) {
+            }
 
-        @Override
-        public void onCompleted() {
-        }
-      });
+            @Override
+            public void onCompleted() {
+            }
+          });
       int counter = 0;
       while (reader.read(root)) {
         stream.putNext(Integer.toString(counter).getBytes(StandardCharsets.UTF_8));