strategies for chunked_arrays, columns, record batches; test the strategies themselves

Author: kszucs

python/pyarrow/tests/strategies.py

@@ -123,8 +123,9 @@ def nested_complex_types(inner_strategy=primitive_types):
     return st.recursive(inner_strategy, complex_types)
 
 
-def schemas(type_strategy=primitive_types):
-    return st.builds(pa.schema, st.lists(fields(type_strategy)))
+def schemas(type_strategy=primitive_types, max_fields=None):
+    children = st.lists(fields(type_strategy), max_size=max_fields)
+    return st.builds(pa.schema, children)
 
 
 complex_schemas = schemas(complex_types())
@@ -135,32 +136,40 @@ def schemas(type_strategy=primitive_types):
 all_schemas = schemas(all_types)
 
 
+_default_array_sizes = st.integers(min_value=0, max_value=20)
+
+
 @st.composite
-def arrays(draw, type, size):
+def arrays(draw, type, size=None):
     if isinstance(type, st.SearchStrategy):
         type = draw(type)
+    elif not isinstance(type, pa.DataType):
+        raise TypeError('Type must be a pyarrow DataType')
+
     if isinstance(size, st.SearchStrategy):
         size = draw(size)
-
-    if not isinstance(type, pa.DataType):
-        raise TypeError('Type must be a pyarrow DataType')
-    if not isinstance(size, int):
+    elif size is None:
+        size = draw(_default_array_sizes)
+    elif not isinstance(size, int):
         raise TypeError('Size must be an integer')
 
     shape = (size,)
 
     if pa.types.is_list(type):
+        # TODO(kszucs) limit the depth
         offsets = draw(npst.arrays(np.uint8(), shape=shape)).cumsum() // 20
         offsets = np.insert(offsets, 0, 0, axis=0)  # prepend with zero
         values = draw(arrays(type.value_type, size=int(offsets.sum())))
         return pa.ListArray.from_arrays(offsets, values)
 
     if pa.types.is_struct(type):
-        h.assume(len(type) > 0)  # TODO issue pa.struct([])
+        h.assume(len(type) > 0)  # TODO(kszucs): create issue -> pa.struct([])
         names, child_arrays = [], []
         for field in type:
             names.append(field.name)
             child_arrays.append(draw(arrays(field.type, size=size)))
+        # fields' metadata are lost here, because from_arrays doesn't accept
+        # a fields argumentum, only names
         return pa.StructArray.from_arrays(child_arrays, names=names)
 
     if (pa.types.is_boolean(type) or pa.types.is_integer(type) or
@@ -181,12 +190,54 @@ def arrays(draw, type, size):
         value = st.binary()
     elif pa.types.is_string(type):
         value = st.text()
-    elif pa.types.is_decimal(type):
-        # FIXME(kszucs): properly limit the precision
-        value = st.decimals(places=type.scale, allow_infinity=False)
-        type = None  # We let arrow infer it from the values
+    # elif pa.types.is_decimal(type):
+    #     # TODO(kszucs): properly limit the precision
+    #     value = st.decimals(places=type.scale, allow_infinity=False)
+    #     type = None  # We let arrow infer it from the values
     else:
+        h.assume(not pa.types.is_decimal(type))
         raise NotImplementedError(type)
 
     values = st.lists(value, min_size=size, max_size=size)
     return pa.array(draw(values), type=type)
+
+
+@st.composite
+def chunked_arrays(draw, type, min_chunks=0, max_chunks=None, chunk_size=None):
+    if isinstance(type, st.SearchStrategy):
+        type = draw(type)
+
+    # TODO(kszucs): remove it, field metadata is not kept
+    h.assume(not pa.types.is_struct(type))
+
+    chunk = arrays(type, size=chunk_size)
+    chunks = st.lists(chunk, min_size=min_chunks, max_size=max_chunks)
+
+    return pa.chunked_array(draw(chunks), type=type)
+
+
+def columns(type, min_chunks=0, max_chunks=None, chunk_size=None):
+    chunked_array = chunked_arrays(type, chunk_size=chunk_size,
+                                   min_chunks=min_chunks,
+                                   max_chunks=max_chunks)
+    return st.builds(pa.column, st.text(), chunked_array)
+
+
+@st.composite
+def record_batches(draw, type, rows=None, max_fields=None):
+    if isinstance(rows, st.SearchStrategy):
+        rows = draw(rows)
+    elif rows is None:
+        rows = draw(_default_array_sizes)
+    elif not isinstance(rows, int):
+        raise TypeError('Rows must be an integer')
+
+    schema = draw(schemas(type, max_fields=max_fields))
+    children = [draw(arrays(field.type, size=rows)) for field in schema]
+    return pa.RecordBatch.from_arrays(children, names=schema)
+
+
+all_arrays = arrays(all_types)
+all_chunked_arrays = chunked_arrays(all_types)
+all_columns = columns(all_types)
+all_record_batches = record_batches(all_types)

python/pyarrow/tests/test_strategies.py

@@ -0,0 +1,49 @@
+import hypothesis as h
+import hypothesis.strategies as st
+
+import pyarrow as pa
+import pyarrow.tests.strategies as past
+
+
+@h.given(past.all_types)
+def test_types(ty):
+    assert isinstance(ty, pa.lib.DataType)
+
+
+@h.given(past.all_fields)
+def test_fields(field):
+    assert isinstance(field, pa.lib.Field)
+
+
+@h.given(past.all_schemas)
+def test_schemas(schema):
+    assert isinstance(schema, pa.lib.Schema)
+
+
+@h.given(past.all_arrays)
+def test_arrays(array):
+    assert isinstance(array, pa.lib.Array)
+
+
+@h.given(past.all_chunked_arrays)
+def test_chunked_arrays(chunked_array):
+    assert isinstance(chunked_array, pa.lib.ChunkedArray)
+
+
+@h.given(past.all_columns)
+def test_columns(column):
+    assert isinstance(column, pa.lib.Column)
+
+
+@h.given(past.all_record_batches)
+def test_record_batches(record_bath):
+    assert isinstance(record_bath, pa.lib.RecordBatch)
+
+
+############################################################
+
+
+@h.given(st.text(), past.all_arrays | past.all_chunked_arrays)
+def test_column_factory(name, arr):
+    column = pa.column(name, arr)
+    assert isinstance(column, pa.Column)