Add max buffer size to Raptor session properties

presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorPageSinkProvider.java

@@ -25,12 +25,12 @@
 import com.facebook.presto.spi.PageSorter;
 import com.facebook.presto.spi.connector.ConnectorPageSinkProvider;
 import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
-import io.airlift.units.DataSize;
 
 import javax.inject.Inject;
 
 import java.util.List;
 
+import static com.facebook.presto.raptor.RaptorSessionProperties.getWriterMaxBufferSize;
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toList;
@@ -41,7 +41,6 @@ public class RaptorPageSinkProvider
     private final StorageManager storageManager;
     private final PageSorter pageSorter;
     private final TemporalFunction temporalFunction;
-    private final DataSize maxBufferSize;
     private final int maxAllowedFilesPerWriter;
 
     @Inject
@@ -50,7 +49,6 @@ public RaptorPageSinkProvider(StorageManager storageManager, PageSorter pageSort
         this.storageManager = requireNonNull(storageManager, "storageManager is null");
         this.pageSorter = requireNonNull(pageSorter, "pageSorter is null");
         this.temporalFunction = requireNonNull(temporalFunction, "temporalFunction is null");
-        this.maxBufferSize = config.getMaxBufferSize();
         this.maxAllowedFilesPerWriter = config.getMaxAllowedFilesPerWriter();
     }
 
@@ -73,7 +71,7 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa
                 handle.getBucketCount(),
                 toColumnIds(handle.getBucketColumnHandles()),
                 handle.getTemporalColumnHandle(),
-                maxBufferSize,
+                getWriterMaxBufferSize(session),
                 maxAllowedFilesPerWriter);
     }
 
@@ -96,7 +94,7 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa
                 handle.getBucketCount(),
                 toColumnIds(handle.getBucketColumnHandles()),
                 handle.getTemporalColumnHandle(),
-                maxBufferSize,
+                getWriterMaxBufferSize(session),
                 maxAllowedFilesPerWriter);
     }
 

presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorSessionProperties.java

@@ -36,6 +36,7 @@ public class RaptorSessionProperties
     private static final String READER_MAX_MERGE_DISTANCE = "reader_max_merge_distance";
     private static final String READER_MAX_READ_SIZE = "reader_max_read_size";
     private static final String READER_STREAM_BUFFER_SIZE = "reader_stream_buffer_size";
+    private static final String WRITER_MAX_BUFFER_SIZE = "writer_max_buffer_size";
     private static final String READER_TINY_STRIPE_THRESHOLD = "reader_tiny_stripe_threshold";
     private static final String READER_LAZY_READ_SMALL_RANGES = "reader_lazy_read_small_ranges";
     private static final String ONE_SPLIT_PER_BUCKET_THRESHOLD = "one_split_per_bucket_threshold";
@@ -86,7 +87,12 @@ public RaptorSessionProperties(StorageManagerConfig config)
                         ORC_ZSTD_JNI_DECOMPRESSION_ENABLED,
                         "use JNI based std decompression for reading ORC files",
                         config.isZstdJniDecompressionEnabled(),
-                        true));
+                        true),
+                dataSizeSessionProperty(
+                        WRITER_MAX_BUFFER_SIZE,
+                        "Raptor page writer max logical buffer size",
+                        config.getMaxBufferSize(),
+                        false));
     }
 
     public List<PropertyMetadata<?>> getSessionProperties()
@@ -119,6 +125,11 @@ public static DataSize getReaderTinyStripeThreshold(ConnectorSession session)
         return session.getProperty(READER_TINY_STRIPE_THRESHOLD, DataSize.class);
     }
 
+    public static DataSize getWriterMaxBufferSize(ConnectorSession session)
+    {
+        return session.getProperty(WRITER_MAX_BUFFER_SIZE, DataSize.class);
+    }
+
     public static boolean isReaderLazyReadSmallRanges(ConnectorSession session)
     {
         return session.getProperty(READER_LAZY_READ_SMALL_RANGES, Boolean.class);

presto-raptor/src/test/java/com/facebook/presto/raptor/TestRaptorSessionProperties.java

@@ -0,0 +1,36 @@
+/*
+ * Licensed 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 com.facebook.presto.raptor;
+
+import com.facebook.presto.raptor.storage.StorageManagerConfig;
+import com.facebook.presto.spi.ConnectorSession;
+import com.facebook.presto.testing.TestingConnectorSession;
+import io.airlift.units.DataSize;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+@Test(singleThreaded = true)
+public class TestRaptorSessionProperties
+{
+    private static final ConnectorSession SESSION = new TestingConnectorSession(
+            new RaptorSessionProperties(new StorageManagerConfig()).getSessionProperties());
+
+    @Test
+    public void testWriterMaxBufferSize()
+    {
+        DataSize maxBufferSize = RaptorSessionProperties.getWriterMaxBufferSize(SESSION);
+        assertEquals(maxBufferSize, new StorageManagerConfig().getMaxBufferSize());
+    }
+}