Add asv benchmarks for plasma

python/CMakeLists.txt

@@ -404,7 +404,7 @@ if (PYARROW_BUILD_PLASMA)
     ${LINK_LIBS}
     libplasma_shared)
 
-  set(CYTHON_EXTENSIONS ${CYTHON_EXTENSIONS} plasma)
+  set(CYTHON_EXTENSIONS ${CYTHON_EXTENSIONS} _plasma)
   file(COPY ${PLASMA_EXECUTABLE} DESTINATION ${BUILD_OUTPUT_ROOT_DIRECTORY})
 endif()
 

python/benchmarks/plasma.py

@@ -0,0 +1,68 @@
+# 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.
+
+import numpy as np
+
+import pyarrow as pa
+import pyarrow.plasma as plasma
+
+from . import common
+
+
+class PlasmaBenchmarks(object):
+
+    def setup(self):
+        self.plasma_store_ctx = plasma.start_plasma_store(plasma_store_memory=10**9)
+        plasma_store_name, p = self.plasma_store_ctx.__enter__()
+        self.plasma_client = plasma.connect(plasma_store_name, "", 64)
+
+        self.data_1kb = np.random.randn(1000 // 8)
+        self.data_10kb = np.random.randn(10000 // 8)
+        self.data_100kb = np.random.randn(100000 // 8)
+        self.data_1mb = np.random.randn(1000000 // 8)
+        self.data_10mb = np.random.randn(10000000 // 8)
+        self.data_100mb = np.random.randn(100000000 // 8)
+
+    def teardown(self):
+        self.plasma_store_ctx.__exit__(None, None, None)
+
+    def time_plasma_put(self):
+        for i in range(1000):
+            self.plasma_client.put(1)
+
+    def time_plasma_putget(self):
+        for i in range(1000):
+            x = self.plasma_client.put(1)
+            self.plasma_client.get(x)
+
+    def time_plasma_put_1kb(self):
+        self.plasma_client.put(self.data_1kb)
+
+    def time_plasma_put_10kb(self):
+        self.plasma_client.put(self.data_10kb)
+
+    def time_plasma_put_100kb(self):
+        self.plasma_client.put(self.data_100kb)
+
+    def time_plasma_put_1mb(self):
+        self.plasma_client.put(self.data_1mb)
+
+    def time_plasma_put_10mb(self):
+        self.plasma_client.put(self.data_10mb)
+
+    def time_plasma_put_100mb(self):
+        self.plasma_client.put(self.data_100mb)

python/pyarrow/plasma.py

@@ -0,0 +1,96 @@
+# 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.
+
+import contextlib
+import os
+import pyarrow as pa
+import subprocess
+import shutil
+import tempfile
+import time
+
+from pyarrow._plasma import *
+
+
+@contextlib.contextmanager
+def start_plasma_store(plasma_store_memory,
+                       use_valgrind=False, use_profiler=False,
+                       use_one_memory_mapped_file=False,
+                       plasma_directory=None, use_hugepages=False):
+    """Start a plasma store process.
+    Args:
+        plasma_store_memory (int): Capacity of the plasma store in bytes.
+        use_valgrind (bool): True if the plasma store should be started inside
+            of valgrind. If this is True, use_profiler must be False.
+        use_profiler (bool): True if the plasma store should be started inside
+            a profiler. If this is True, use_valgrind must be False.
+        use_one_memory_mapped_file: If True, then the store will use only a
+            single memory-mapped file.
+        plasma_directory (str): Directory where plasma memory mapped files
+            will be stored.
+        use_hugepages (bool): True if the plasma store should use huge pages.
+    Return:
+        A tuple of the name of the plasma store socket and the process ID of
+            the plasma store process.
+    """
+    if use_valgrind and use_profiler:
+        raise Exception("Cannot use valgrind and profiler at the same time.")
+
+    tmpdir = tempfile.mkdtemp(prefix='test_plasma-')
+    try:
+        plasma_store_name = os.path.join(tmpdir, 'plasma.sock')
+        plasma_store_executable = os.path.join(pa.__path__[0], "plasma_store")
+        command = [plasma_store_executable,
+                   "-s", plasma_store_name,
+                   "-m", str(plasma_store_memory)]
+        if use_one_memory_mapped_file:
+            command += ["-f"]
+        if plasma_directory:
+            command += ["-d", plasma_directory]
+        if use_hugepages:
+            command += ["-h"]
+        stdout_file = None
+        stderr_file = None
+        if use_valgrind:
+            command = ["valgrind",
+                       "--track-origins=yes",
+                       "--leak-check=full",
+                       "--show-leak-kinds=all",
+                       "--leak-check-heuristics=stdstring",
+                       "--error-exitcode=1"] + command
+            proc = subprocess.Popen(command, stdout=stdout_file,
+                                    stderr=stderr_file)
+            time.sleep(1.0)
+        elif use_profiler:
+            command = ["valgrind", "--tool=callgrind"] + command
+            proc = subprocess.Popen(command, stdout=stdout_file,
+                                    stderr=stderr_file)
+            time.sleep(1.0)
+        else:
+            proc = subprocess.Popen(command, stdout=stdout_file,
+                                    stderr=stderr_file)
+            time.sleep(0.1)
+        rc = proc.poll()
+        if rc is not None:
+            raise RuntimeError("plasma_store exited unexpectedly with "
+                               "code %d" % (rc,))
+
+        yield plasma_store_name, proc
+    finally:
+        if proc.poll() is None:
+            proc.kill()
+        shutil.rmtree(tmpdir)

python/pyarrow/tests/test_plasma.py

@@ -106,76 +106,6 @@ def assert_get_object_equal(unit_test, client1, client2, object_id,
         assert plasma.buffers_equal(metadata, client1_metadata)
 
 
-@contextlib.contextmanager
-def start_plasma_store(plasma_store_memory=DEFAULT_PLASMA_STORE_MEMORY,
-                       use_valgrind=False, use_profiler=False,
-                       use_one_memory_mapped_file=False,
-                       plasma_directory=None, use_hugepages=False):
-    """Start a plasma store process.
-    Args:
-        use_valgrind (bool): True if the plasma store should be started inside
-            of valgrind. If this is True, use_profiler must be False.
-        use_profiler (bool): True if the plasma store should be started inside
-            a profiler. If this is True, use_valgrind must be False.
-        stdout_file: A file handle opened for writing to redirect stdout to. If
-            no redirection should happen, then this should be None.
-        stderr_file: A file handle opened for writing to redirect stderr to. If
-            no redirection should happen, then this should be None.
-        use_one_memory_mapped_file: If True, then the store will use only a
-            single memory-mapped file.
-    Return:
-        A tuple of the name of the plasma store socket and the process ID of
-            the plasma store process.
-    """
-    if use_valgrind and use_profiler:
-        raise Exception("Cannot use valgrind and profiler at the same time.")
-
-    tmpdir = tempfile.mkdtemp(prefix='test_plasma-')
-    try:
-        plasma_store_name = os.path.join(tmpdir, 'plasma.sock')
-        plasma_store_executable = os.path.join(pa.__path__[0], "plasma_store")
-        command = [plasma_store_executable,
-                   "-s", plasma_store_name,
-                   "-m", str(plasma_store_memory)]
-        if use_one_memory_mapped_file:
-            command += ["-f"]
-        if plasma_directory:
-            command += ["-d", plasma_directory]
-        if use_hugepages:
-            command += ["-h"]
-        stdout_file = None
-        stderr_file = None
-        if use_valgrind:
-            command = ["valgrind",
-                       "--track-origins=yes",
-                       "--leak-check=full",
-                       "--show-leak-kinds=all",
-                       "--leak-check-heuristics=stdstring",
-                       "--error-exitcode=1"] + command
-            proc = subprocess.Popen(command, stdout=stdout_file,
-                                    stderr=stderr_file)
-            time.sleep(1.0)
-        elif use_profiler:
-            command = ["valgrind", "--tool=callgrind"] + command
-            proc = subprocess.Popen(command, stdout=stdout_file,
-                                    stderr=stderr_file)
-            time.sleep(1.0)
-        else:
-            proc = subprocess.Popen(command, stdout=stdout_file,
-                                    stderr=stderr_file)
-            time.sleep(0.1)
-        rc = proc.poll()
-        if rc is not None:
-            raise RuntimeError("plasma_store exited unexpectedly with "
-                               "code %d" % (rc,))
-
-        yield plasma_store_name, proc
-    finally:
-        if proc.poll() is None:
-            proc.kill()
-        shutil.rmtree(tmpdir)
-
-
 @pytest.mark.plasma
 class TestPlasmaClient(object):
 
@@ -185,7 +115,8 @@ class TestPlasmaClient(object):
 
         import pyarrow.plasma as plasma
         # Start Plasma store.
-        self.plasma_store_ctx = start_plasma_store(
+        self.plasma_store_ctx = plasma.start_plasma_store(
+            plasma_store_memory=DEFAULT_PLASMA_STORE_MEMORY,
             use_valgrind=USE_VALGRIND,
             use_one_memory_mapped_file=use_one_memory_mapped_file)
         plasma_store_name, self.p = self.plasma_store_ctx.__enter__()
@@ -839,8 +770,10 @@ def test_object_id_equality_operators():
                     reason="requires hugepage support")
 def test_use_huge_pages():
     import pyarrow.plasma as plasma
-    with start_plasma_store(plasma_directory="/mnt/hugepages",
-                            use_hugepages=True) as (plasma_store_name, p):
+    with plasma.start_plasma_store(
+            plasma_store_memory=DEFAULT_PLASMA_STORE_MEMORY,
+            plasma_directory="/mnt/hugepages",
+            use_hugepages=True) as (plasma_store_name, p):
         plasma_client = plasma.connect(plasma_store_name, "", 64)
         create_object(plasma_client, 100000000)
 
@@ -852,7 +785,9 @@ def test_use_huge_pages():
 def test_plasma_client_sharing():
     import pyarrow.plasma as plasma
 
-    with start_plasma_store() as (plasma_store_name, p):
+    with plasma.start_plasma_store(
+            plasma_store_memory=DEFAULT_PLASMA_STORE_MEMORY) \
+                as (plasma_store_name, p):
         plasma_client = plasma.connect(plasma_store_name, "", 64)
         object_id = plasma_client.put(np.zeros(3))
         buf = plasma_client.get(object_id)

python/setup.py

@@ -144,7 +144,7 @@ class build_ext(_build_ext):
         'lib',
         '_parquet',
         '_orc',
-        'plasma']
+        '_plasma']
 
     def _run_cmake(self):
         # The directory containing this setup.py