if sys.platform in ('win32', 'darwin'):

pytest.skip("need sparse file support")

try:

filename = tempfile.mktemp(prefix='test_io')

with open(filename, 'wb') as f:

f.truncate(2 ** 32 + 10)

f.seek(2 ** 32 + 5)

f.write(b'mark\n')

with file_factory(filename) as f:

assert f.seek(2 ** 32 + 5) == 2 ** 32 + 5

assert f.tell() == 2 ** 32 + 5

assert f.read(5) == b'mark\n'

assert f.tell() == 2 ** 32 + 10

finally:

buf = BytesIO()

f = pa.PythonFile(buf)

assert f.tell() == 0

s1 = b'enga\xc3\xb1ado'

s2 = b'foobar'

f.write(s1)

assert f.tell() == len(s1)

f.write(s2)

expected = s1 + s2

result = buf.getvalue()

assert result == expected

assert not f.closed

f.close()

assert f.closed

with pytest.raises(TypeError, match="binary file expected"):

data = b'some sample data'

buf = BytesIO(data)

f = pa.PythonFile(buf, mode='r')

assert f.size() == len(data)

assert f.tell() == 0

assert f.read(4) == b'some'

assert f.tell() == 4

f.seek(0)

assert f.tell() == 0

f.seek(5)

assert f.tell() == 5

v = f.read(50)

assert v == b'sample data'

assert len(v) == 11

assert f.size() == len(data)

assert not f.closed

f.close()

assert f.closed

with pytest.raises(TypeError, match="binary file expected"):

data = b'data1data2data3data4data5'

f = pa.PythonFile(BytesIO(data), mode='r')

# negative nbytes or offsets don't make sense here, raise ValueError

if nbytes < 0 or file_offset < 0:

with pytest.raises(pa.ArrowInvalid,

match="should be a positive value"):

f.get_stream(file_offset=file_offset, nbytes=nbytes)

f.close()

return

else:

stream = f.get_stream(file_offset=file_offset, nbytes=nbytes)

# Subsequent calls to 'read' should match behavior if same

# data passed to BytesIO where get_stream should handle if

# nbytes/file_offset results in no bytes b/c out of bounds.

start = min(file_offset, len(data))

end = min(file_offset + nbytes, len(data))

buf = BytesIO(data[start:end])

# read some chunks

assert stream.read(nbytes=4) == buf.read(4)

assert stream.read(nbytes=6) == buf.read(6)

# Read to end of each stream

assert stream.read() == buf.read()

# Try reading past the stream

n = len(data) * 2

assert stream.read(n) == buf.read(n)

# NativeFile[CInputStream] is not seekable

with pytest.raises(OSError, match="seekable"):

stream.seek(0)

stream.close()

data = b'some sample data'

buf = BytesIO(data)

f = pa.PythonFile(buf, mode='r')

# test simple read at

v = f.read_at(nbytes=5, offset=3)

assert v == b'e sam'

assert len(v) == 5

# test reading entire file when nbytes > len(file)

w = f.read_at(nbytes=50, offset=0)

assert w == data

data = b'some sample data'

buf = BytesIO(data)

with pa.PythonFile(buf, mode='r') as f:

length = 10

data = b'some sample data longer than 10'

dst_buf = bytearray(length)

src_buf = BytesIO(data)

with pa.PythonFile(src_buf, mode='r') as f:

assert f.readinto(dst_buf) == 10

assert dst_buf[:length] == data[:length]

length = 10

data = b'0123456798'

dst_buf = bytearray(data)

class DuckReader:

def close(self):

pass

@property

def closed(self):

return False

def read_buffer(self, nbytes):

assert nbytes == length

return memoryview(dst_buf)[:nbytes]

duck_reader = DuckReader()

with pa.PythonFile(duck_reader, mode='r') as f:

buf = f.read_buffer(length)

assert len(buf) == length

assert memoryview(buf).tobytes() == dst_buf[:length]

# buf should point to the same memory, so modyfing it

memoryview(buf)[0] = ord(b'x')

# should modify the original

data = b'''line1

buf = BytesIO(data)

buf2 = BytesIO(data)

with pa.PythonFile(buf, mode='r') as f:

for read, expected in zip(f, buf2):

def factory(filename):

return pa.PythonFile(open(filename, 'rb'))

# Like a BytesIO, but zero-copy underneath for C++ consumers

data = b'some sample data'

f = pa.BufferReader(data)

assert f.tell() == 0

assert f.size() == len(data)

assert f.read(4) == b'some'

assert f.tell() == 4

f.seek(0)

assert f.tell() == 0

f.seek(0, 2)

assert f.tell() == len(data)

f.seek(5)

assert f.tell() == 5

assert f.read(50) == b'sample data'

assert not f.closed

f.close()

import gc

# ARROW-421

def get_buffer():

data = b'some sample data' * 1000

reader = pa.BufferReader(data)

reader.seek(5)

return reader.read_buffer(6)

buf = get_buffer()

gc.collect()

assert buf.to_pybytes() == b'sample'

path = os.path.join(str(tmpdir), 'foo.txt')

with open(path, 'wb') as f:

pf = pa.PythonFile(f)

assert pf.writable()

assert not pf.readable()

assert not pf.seekable() # PyOutputStream isn't seekable

f.write(b'foobar\n')

with open(path, 'rb') as f:

pf = pa.PythonFile(f)

assert pf.readable()

assert not pf.writable()

assert pf.seekable()

assert pf.read() == b'foobar\n'

bio = BytesIO()

pf = pa.PythonFile(bio)

assert pf.writable()

assert not pf.readable()

assert not pf.seekable()

pf.write(b'foobar\n')

lines = [b'line1\n', b'line2\n' b'line3']

path = os.path.join(str(tmpdir), 'foo.txt')

with open(path, 'wb') as f:

try:

f = pa.PythonFile(f, mode='w')

assert f.writable()

f.writelines(lines)

finally:

f.close()

with open(path, 'rb') as f:

try:

f = pa.PythonFile(f, mode='r')

assert f.readable()

assert f.read() == b''.join(lines)

finally:

bio = BytesIO()

pf = pa.PythonFile(bio)

wr = weakref.ref(pf)

del pf

assert wr() is None # object was destroyed

assert not bio.closed

pf = pa.PythonFile(bio)

pf.close()

# Check that buffer survives a pickle roundtrip

for protocol in range(0, pickle.HIGHEST_PROTOCOL + 1):

result = pickle.loads(pickle.dumps(buf, protocol=protocol))

assert len(result) == len(buf)

assert memoryview(result) == memoryview(buf)

assert result.to_pybytes() == buf.to_pybytes()

val = b'some data'

buf = pa.py_buffer(val)

assert isinstance(buf, pa.Buffer)

assert not buf.is_mutable

assert buf.is_cpu

result = buf.to_pybytes()

assert result == val

null_buff = pa.foreign_buffer(address=0, size=0)

assert null_buff.to_pybytes() == b""

assert null_buff.address == 0

# ARROW-16048: we shouldn't expose a NULL address through the Python

# buffer protocol.

m = memoryview(null_buff)

assert m.tobytes() == b""

assert pa.py_buffer(m).address != 0

val = b'some data'

buf = pa.py_buffer(val)

assert isinstance(buf, pa.Buffer)

assert not buf.is_mutable

assert buf.is_cpu

result = memoryview(buf)

assert result == val

val = bytearray(b'some data')

buf = pa.py_buffer(val)

assert isinstance(buf, pa.Buffer)

assert buf.is_mutable

assert buf.is_cpu

result = bytearray(buf)

assert result == val

buf = pa.py_buffer(b'some data')

wr = weakref.ref(buf)

assert wr() is not None

del buf

# Make sure creating a numpy array from an arrow buffer works

byte_array = bytearray(20)

byte_array[0] = 42

buf = pa.py_buffer(byte_array)

array = np.frombuffer(buf, dtype="uint8")

assert array[0] == byte_array[0]

byte_array[0] += 1

assert array[0] == byte_array[0]

# C-contiguous

arr = np.arange(12, dtype=np.int8).reshape((3, 4))

buf = pa.py_buffer(arr)

assert buf.is_cpu

assert buf.is_mutable

assert buf.to_pybytes() == arr.tobytes()

# F-contiguous; note strides information is lost

buf = pa.py_buffer(arr.T)

assert buf.is_cpu

assert buf.is_mutable

assert buf.to_pybytes() == arr.tobytes()

# Non-contiguous

with pytest.raises(ValueError, match="not contiguous"):

b1 = b'some data!'

b2 = bytearray(b1)

b3 = bytearray(b1)

buf1 = pa.py_buffer(b1)

buf2 = pa.py_buffer(b1)

buf3 = pa.py_buffer(b2)

buf4 = pa.py_buffer(b3)

assert buf1.address > 0

assert buf1.address == buf2.address

assert buf3.address != buf2.address

assert buf4.address != buf3.address

arr = np.arange(5)

buf = pa.py_buffer(arr)

# Buffer.equals() returns true iff the buffers have the same contents

def eq(a, b):

assert a.equals(b)

assert a == b

assert not (a != b)

def ne(a, b):

assert not a.equals(b)

assert not (a == b)

assert a != b

b1 = b'some data!'

b2 = bytearray(b1)

b3 = bytearray(b1)

b3[0] = 42

buf1 = pa.py_buffer(b1)

buf2 = pa.py_buffer(b2)

buf3 = pa.py_buffer(b2)

buf4 = pa.py_buffer(b3)

buf5 = pa.py_buffer(np.frombuffer(b2, dtype=np.int16))

eq(buf1, buf1)

eq(buf1, buf2)

eq(buf2, buf3)

ne(buf2, buf4)

# Data type is indifferent

buf = pa.py_buffer(b'some data')

assert buf == b'some data'

assert buf == bytearray(b'some data')

assert buf != b'some dat1'

with pytest.raises(TypeError):

data = bytearray(b'some data!')

buf = pa.py_buffer(data)

n = len(data)

for ix in range(-n, n - 1):

assert buf[ix] == data[ix]

with pytest.raises(IndexError):

buf[n]

with pytest.raises(IndexError):

data = b'some data!'

buf = pa.py_buffer(data)

sliced = buf.slice(2)

expected = pa.py_buffer(b'me data!')

assert sliced.equals(expected)

sliced2 = buf.slice(2, 4)

expected2 = pa.py_buffer(b'me d')

assert sliced2.equals(expected2)

# 0 offset

assert buf.slice(0).equals(buf)

# Slice past end of buffer

assert len(buf.slice(len(buf))) == 0

with pytest.raises(IndexError):

buf.slice(-1)

with pytest.raises(IndexError):

buf.slice(len(buf) + 1)

assert buf[11:].to_pybytes() == b""

# Slice stop exceeds buffer length

with pytest.raises(IndexError):

buf.slice(1, len(buf))

assert buf[1:11].to_pybytes() == buf.to_pybytes()[1:]

# Negative length

with pytest.raises(IndexError):

buf.slice(1, -1)

# Test slice notation

assert buf[2:].equals(buf.slice(2))

assert buf[2:5].equals(buf.slice(2, 3))

assert buf[-5:].equals(buf.slice(len(buf) - 5))

assert buf[-5:-2].equals(buf.slice(len(buf) - 5, 3))

with pytest.raises(IndexError):

buf[::-1]

with pytest.raises(IndexError):

buf[::2]

n = len(buf)

for start in range(-n * 2, n * 2):

for stop in range(-n * 2, n * 2):

# ARROW-3228; NumPy's frombuffer ctor determines whether a buffer-like

# object is mutable by first attempting to get a mutable buffer using

# PyObject_FromBuffer. If that fails, it assumes that the object is

# immutable

a = b'12345'

arrow_ref = pa.py_buffer(a)

numpy_ref = np.frombuffer(arrow_ref, dtype=np.uint8)

obj = np.array([1, 2], dtype=np.int32)

addr = obj.__array_interface__["data"][0]

size = obj.nbytes

buf = pa.foreign_buffer(addr, size, obj)

wr = weakref.ref(obj)

del obj

assert np.frombuffer(buf, dtype=np.int32).tolist() == [1, 2]

assert wr() is not None

del buf

buf = pa.allocate_buffer(100)

assert buf.size == 100

assert buf.is_mutable

assert buf.parent is None

bit = b'abcde'

writer = pa.FixedSizeBufferWriter(buf)

writer.write(bit)

buf = pa.allocate_buffer(100, resizable=True)

assert isinstance(buf, pa.ResizableBuffer)

buf.resize(200)

if not Codec.is_available(compression):

pytest.skip("{} support is not built".format(compression))

INPUT_SIZE = 10000

test_data = (np.random.randint(0, 255, size=INPUT_SIZE)

.astype(np.uint8)

.tobytes())

test_buf = pa.py_buffer(test_data)

compressed_buf = pa.compress(test_buf, codec=compression)

compressed_bytes = pa.compress(test_data, codec=compression,

asbytes=True)

assert isinstance(compressed_bytes, bytes)

decompressed_buf = pa.decompress(compressed_buf, INPUT_SIZE,

codec=compression)

decompressed_bytes = pa.decompress(compressed_bytes, INPUT_SIZE,

codec=compression, asbytes=True)

assert isinstance(decompressed_bytes, bytes)

assert decompressed_buf.equals(test_buf)

assert decompressed_bytes == test_data

with pytest.raises(ValueError):

if not Codec.is_available(compression):

pytest.skip("{} support is not built".format(compression))

codec = Codec(compression)

if codec.name == "snappy":

assert codec.compression_level is None

else:

assert isinstance(codec.compression_level, int)

# These codecs do not support a compression level

no_level = ['snappy']

if compression in no_level:

assert not Codec.supports_compression_level(compression)

with pytest.raises(ValueError):

Codec(compression, 0)

with pytest.raises(ValueError):

Codec.minimum_compression_level(compression)

with pytest.raises(ValueError):

Codec.maximum_compression_level(compression)

with pytest.raises(ValueError):

Codec.default_compression_level(compression)

return

INPUT_SIZE = 10000

test_data = (np.random.randint(0, 255, size=INPUT_SIZE)

.astype(np.uint8)

.tobytes())

test_buf = pa.py_buffer(test_data)

min_level = Codec.minimum_compression_level(compression)

max_level = Codec.maximum_compression_level(compression)

default_level = Codec.default_compression_level(compression)

assert min_level < max_level

assert default_level >= min_level

assert default_level <= max_level

for compression_level in range(min_level, max_level+1):

codec = Codec(compression, compression_level)

compressed_buf = codec.compress(test_buf)

compressed_bytes = codec.compress(test_data, asbytes=True)

assert isinstance(compressed_bytes, bytes)

decompressed_buf = codec.decompress(compressed_buf, INPUT_SIZE)

decompressed_bytes = codec.decompress(compressed_bytes, INPUT_SIZE,

asbytes=True)

assert isinstance(decompressed_bytes, bytes)

assert decompressed_buf.equals(test_buf)

assert decompressed_bytes == test_data

with pytest.raises(ValueError):

codec.decompress(compressed_bytes)

# The ability to set a seed this way is not present on older versions of

# numpy (currently in our python 3.6 CI build). Some inputs might just

# happen to compress the same between the two levels so using seeded

# random numbers is necessary to help get more reliable results

#

# The goal of this part is to ensure the compression_level is being

# passed down to the C++ layer, not to verify the compression algs

# themselves

if not hasattr(np.random, 'default_rng'):

pytest.skip('Requires newer version of numpy')

rng = np.random.default_rng(seed=42)

values = rng.integers(0, 100, 1000)

arr = pa.array(values)

hard_to_compress_buffer = arr.buffers()[1]

weak_codec = Codec(compression, min_level)

weakly_compressed_buf = weak_codec.compress(hard_to_compress_buffer)

strong_codec = Codec(compression, max_level)

strongly_compressed_buf = strong_codec.compress(hard_to_compress_buffer)

val = b'some data'

buf = pa.py_buffer(val)

assert not buf.is_mutable

assert isinstance(buf, pa.Buffer)

result = memoryview(buf)

assert result.readonly

with pytest.raises(TypeError) as exc:

result[0] = b'h'

assert 'cannot modify read-only' in str(exc.value)

b = bytes(buf)

with pytest.raises(TypeError) as exc:

b[0] = b'h'

# ARROW-2039: calling Buffer() directly creates an uninitialized object

# ARROW-2638: prevent calling extension class constructors directly

with pytest.raises(TypeError):

# 10 bytes

val = b'dataabcdef'

f = pa.BufferOutputStream()

K = 1000

for i in range(K):

f.write(val)

buf = f.getvalue()

assert len(buf) == len(val) * K

f = pa.BufferOutputStream()

f.write(b'ok')

assert not f.closed

f.getvalue()

assert f.closed

with pytest.raises(ValueError):

def make_buffer(bytes_obj):

return bytearray(pa.py_buffer(bytes_obj))

buf = make_buffer(b'foo')

gc.collect()

assert buf == b'foo'

# ARROW-1053

val = b'foo'

refcount_before = sys.getrefcount(val)

for i in range(10):

make_buffer(val)

gc.collect()

f = pa.BufferOutputStream()

data = b'ok'

arr = np.frombuffer(data, dtype='S1')

f.write(arr)

f.write(bytearray(data))

f.write(pa.py_buffer(data))

with pytest.raises(TypeError):

f.write(data.decode('utf8'))

buf = f.getvalue()

# Make sure that the MockOutputStream and the BufferOutputStream record the

# same size

# 10 bytes

val = b'dataabcdef'

f1 = pa.MockOutputStream()

f2 = pa.BufferOutputStream()

K = 1000

for i in range(K):

f1.write(val)

f2.write(val)

assert f1.size() == len(f2.getvalue())

# Do the same test with a table

record_batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], ['a'])

f1 = pa.MockOutputStream()

f2 = pa.BufferOutputStream()

stream_writer1 = pa.RecordBatchStreamWriter(f1, record_batch.schema)

stream_writer2 = pa.RecordBatchStreamWriter(f2, record_batch.schema)

stream_writer1.write_batch(record_batch)

stream_writer2.write_batch(record_batch)

stream_writer1.close()

stream_writer2.close()

SIZE = 4096

arr = np.random.randint(0, 256, size=SIZE).astype('u1')

data = arr.tobytes()[:SIZE]

path = os.path.join(str(tmpdir), guid())

with open(path, 'wb') as f:

f.write(data)

def teardown():

_try_delete(path)

request.addfinalizer(teardown)

allow_read_out_of_bounds=True):

path, data = sample_data

f = FACTORY(path, mode='r')

assert f.read(10) == data[:10]

assert f.read(0) == b''

assert f.tell() == 10

assert f.read() == data[10:]

assert f.size() == len(data)

f.seek(0)

assert f.tell() == 0

# Seeking past end of file not supported in memory maps

if allow_read_out_of_bounds:

f.seek(len(data) + 1)

assert f.tell() == len(data) + 1

assert f.read(5) == b''

# Test whence argument of seek, ARROW-1287

assert f.seek(3) == 3

assert f.seek(3, os.SEEK_CUR) == 6

assert f.tell() == 6

ex_length = len(data) - 2

assert f.seek(-2, os.SEEK_END) == ex_length

path, data = sample_disk_data

cases = []

with pa.memory_map(path, 'rb') as f:

cases.append((f.read_buffer(100), data[:100]))

cases.append((f.read_buffer(100), data[100:200]))

cases.append((f.read_buffer(100), data[200:300]))

# Call gc.collect() for good measure

gc.collect()

for buf, expected in cases: