ARROW-6926: [Python] Support __sizeof__ protocol for Python objects

jorisvandenbossche · pitrou · commit c6bec1575059 · 2019-12-02T14:44:45.000+01:00
https://issues.apache.org/jira/browse/ARROW-6926 Closes apache#5879 from jorisvandenbossche/ARROW-6926-sizeof and squashes the following commits: baca02f <Joris Van den Bossche> check size of metadata keys and values 17f75d4 <Joris Van den Bossche> add Schema sizeof 54896e1 <Joris Van den Bossche> object -> super b36aa5d <Joris Van den Bossche> ARROW-6926: Support __sizeof__ protocol for Python objects Authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com> Signed-off-by: Antoine Pitrou <antoine@python.org>
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
@@ -805,6 +805,9 @@ cdef class Array(_PandasConvertible):
                 size += buf.size
         return size
 
+    def __sizeof__(self):
+        return super(Array, self).__sizeof__() + self.nbytes
+
     def __iter__(self):
         for i in range(len(self)):
             yield self.getitem(i)
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
@@ -125,6 +125,9 @@ cdef class ChunkedArray(_PandasConvertible):
             size += chunk.nbytes
         return size
 
+    def __sizeof__(self):
+        return super(ChunkedArray, self).__sizeof__() + self.nbytes
+
     def __iter__(self):
         for chunk in self.iterchunks():
             for item in chunk:
@@ -657,6 +660,9 @@ cdef class RecordBatch(_PandasConvertible):
             size += self.column(i).nbytes
         return size
 
+    def __sizeof__(self):
+        return super(RecordBatch, self).__sizeof__() + self.nbytes
+
     def __getitem__(self, key):
         if isinstance(key, slice):
             return _normalize_slice(self, key)
@@ -1521,6 +1527,9 @@ cdef class Table(_PandasConvertible):
             size += column.nbytes
         return size
 
+    def __sizeof__(self):
+        return super(Table, self).__sizeof__() + self.nbytes
+
     def add_column(self, int i, field_, column):
         """
         Add column to Table at position. Returns new table
diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py
@@ -1607,13 +1607,16 @@ def test_buffers_nested():
     assert struct.unpack('4xh', values) == (43,)
 
 
-def test_nbytes():
+def test_nbytes_sizeof():
     a = pa.array(np.array([4, 5, 6], dtype='int64'))
     assert a.nbytes == 8 * 3
+    assert sys.getsizeof(a) >= object.__sizeof__(a) + a.nbytes
     a = pa.array([1, None, 3], type='int64')
     assert a.nbytes == 8*3 + 1
+    assert sys.getsizeof(a) >= object.__sizeof__(a) + a.nbytes
     a = pa.array([[1, 2], None, [3, None, 4, 5]], type=pa.list_(pa.int64()))
     assert a.nbytes == 1 + 4 * 4 + 1 + 6 * 8
+    assert sys.getsizeof(a) >= object.__sizeof__(a) + a.nbytes
 
 
 def test_invalid_tensor_constructor_repr():
diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py
@@ -17,6 +17,7 @@
 
 from collections import OrderedDict
 import pickle
+import sys
 
 import pytest
 import numpy as np
@@ -542,3 +543,17 @@ def test_schema_from_pandas():
         schema = pa.Schema.from_pandas(df)
         expected = pa.Table.from_pandas(df).schema
         assert schema == expected
+
+
+def test_schema_sizeof():
+    schema = pa.schema([
+        pa.field('foo', pa.int32()),
+        pa.field('bar', pa.string()),
+    ])
+
+    assert sys.getsizeof(schema) > 30
+
+    schema2 = schema.with_metadata({"key": "some metadata"})
+    assert sys.getsizeof(schema2) > sys.getsizeof(schema)
+    schema3 = schema.with_metadata({"key": "some more metadata"})
+    assert sys.getsizeof(schema3) > sys.getsizeof(schema2)
diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py
@@ -45,6 +45,7 @@ def test_chunked_array_basics():
     assert all(isinstance(c, pa.lib.Int64Array) for c in data.iterchunks())
     assert len(data.chunks) == 3
     assert data.nbytes == sum(c.nbytes for c in data.iterchunks())
+    assert sys.getsizeof(data) >= object.__sizeof__(data) + data.nbytes
     data.validate()
 
 
@@ -276,6 +277,7 @@ def test_recordbatch_basics():
     assert batch.num_rows == 5
     assert batch.num_columns == len(data)
     assert batch.nbytes == 5 * 2 + 1 + 5 * 4 + 1
+    assert sys.getsizeof(batch) >= object.__sizeof__(batch) + batch.nbytes
     pydict = batch.to_pydict()
     assert pydict == OrderedDict([
         ('c0', [0, 1, 2, 3, 4]),
@@ -524,6 +526,7 @@ def test_table_basics():
     assert table.num_columns == 2
     assert table.shape == (5, 2)
     assert table.nbytes == 2 * (5 * 8 + 1)
+    assert sys.getsizeof(table) >= object.__sizeof__(table) + table.nbytes
     pydict = table.to_pydict()
     assert pydict == OrderedDict([
         ('a', [0, 1, 2, 3, 4]),
diff --git a/python/pyarrow/types.pxi b/python/pyarrow/types.pxi
@@ -17,6 +17,7 @@
 
 import atexit
 import re
+import sys
 import warnings
 
 from pyarrow import compat
@@ -891,6 +892,15 @@ cdef class Schema:
     def __hash__(self):
         return hash((tuple(self), self.metadata))
 
+    def __sizeof__(self):
+        size = 0
+        if self.metadata:
+            for key, value in self.metadata.items():
+                size += sys.getsizeof(key)
+                size += sys.getsizeof(value)
+
+        return size + super(Schema, self).__sizeof__()
+
     @property
     def pandas_metadata(self):
         """
