29579092494

Committed 17 Jul 2026 12:06PM UTC coverage: 96.781% (+0.04%) from 96.745%

Build # 29579092494

Build Type

Pull #509

github

Committed by

web-flow

Commit Message

Merge 08de70021 into d056c9e06

Pull Request Pull Request #509: PyArrow support

Coverage Stats

429 of 443 new or added lines in 15 files covered. (96.84%)

7 existing lines in 3 files now uncovered.

5171 of 5343 relevant lines covered (96.78%)

166.23 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.55

/clickhouse_driver/columns/arrow/stringcolumn.py

import numpy as np

from ..base import CommonSerialization
from ..numpy.stringcolumn import (
    NumpyByteFixedString, NumpyByteStringColumn, NumpyFixedString,
    NumpyStringColumn
)
from .base import ArrowColumnMixin


class ArrowStringBuffers(object):
    """
    String column read into Arrow-style buffers: concatenated bytes
    plus int64 offsets, no per-string Python objects. Assembled into
    a pyarrow array by ``clickhouse_driver.arrow.convert``.
    """
    __slots__ = ('offsets', 'data', 'nulls_map')

    def __init__(self, offsets, data, nulls_map=None):
        self.offsets = offsets
        self.data = data
        self.nulls_map = nulls_map

    def __len__(self):
        return len(self.offsets) // 8 - 1


class ArrowStringMixin(ArrowColumnMixin):
    # Reset for columns whose consumers need real items, e.g.
    # LowCardinality dictionaries.
    arrow_buffers_allowed = True

    def _buffers_encoding_ok(self):
        return self.encoding.lower() in ('utf-8', 'utf8')

    def _use_arrow_buffers(self, buf):
        return (
            self.arrow_buffers_allowed and
            self._buffers_encoding_ok() and
            type(self.serialization) is CommonSerialization and
            hasattr(buf, 'read_strings_arrow')
        )

    def _read_data(self, n_items, buf, nulls_map=None):
        if self._use_arrow_buffers(buf):
            offsets, data = buf.read_strings_arrow(n_items)
            return ArrowStringBuffers(offsets, data, nulls_map=nulls_map)

        return super(ArrowStringMixin, self)._read_data(
            n_items, buf, nulls_map=nulls_map
        )

    def _wrap_items(self, items):
        # Wrapping strings into an ndarray re-encodes them (unicode
        # dtype). Arrow consumes the raw tuple directly; the ndarray
        # is only needed for the nullable masked path.
        if self.nullable:
            return np.array(items, dtype=self.dtype)
        return items


class ArrowStringColumn(ArrowStringMixin, NumpyStringColumn):
    def read_items(self, n_items, buf):
        return self._wrap_items(
            buf.read_strings(n_items, encoding=self.encoding)
        )


class ArrowByteStringColumn(ArrowStringMixin, NumpyByteStringColumn):
    def _buffers_encoding_ok(self):
        # No decoding happens for byte strings: buffers can be used
        # regardless of encoding.
        return True

    def read_items(self, n_items, buf):
        return self._wrap_items(buf.read_strings(n_items))


class ArrowFixedString(ArrowStringMixin, NumpyFixedString):
    def _use_arrow_buffers(self, buf):
        # Fixed strings are trimmed of zero bytes on read: buffers
        # would keep the padding.
        return False

    def read_items(self, n_items, buf):
        return self._wrap_items(buf.read_fixed_strings(
            n_items, self.length, encoding=self.encoding
        ))


class ArrowByteFixedString(ArrowStringMixin, NumpyByteFixedString):
    def _use_arrow_buffers(self, buf):
        return False

    def read_items(self, n_items, buf):
        return self._wrap_items(
            buf.read_fixed_strings(n_items, self.length)
        )


def create_arrow_string_column(spec, column_options):
    client_settings = column_options['context'].client_settings
    strings_as_bytes = client_settings['strings_as_bytes']
    encoding = client_settings.get(
        'strings_encoding', NumpyStringColumn.default_encoding
    )

    if spec == 'String':
        cls = ArrowByteStringColumn if strings_as_bytes \
            else ArrowStringColumn
        return cls(encoding=encoding, **column_options)
    else:
        length = int(spec[12:-1])
        cls = ArrowByteFixedString if strings_as_bytes \
            else ArrowFixedString
        return cls(length, encoding=encoding, **column_options)

1	import numpy as np	3✔
2
3	from ..base import CommonSerialization	2✔
4	from ..numpy.stringcolumn import (	2✔
5	NumpyByteFixedString, NumpyByteStringColumn, NumpyFixedString,
6	NumpyStringColumn
7	)
8	from .base import ArrowColumnMixin	2✔
9
10
11	class ArrowStringBuffers(object):	2✔
12	"""
13	String column read into Arrow-style buffers: concatenated bytes
14	plus int64 offsets, no per-string Python objects. Assembled into
15	a pyarrow array by ``clickhouse_driver.arrow.convert``.
16	"""
17	__slots__ = ('offsets', 'data', 'nulls_map')	2✔
18
19	def __init__(self, offsets, data, nulls_map=None):	2✔
20	self.offsets = offsets	2✔
21	self.data = data	2✔
22	self.nulls_map = nulls_map	2✔
23
24	def __len__(self):	2✔
25	return len(self.offsets) // 8 - 1	2✔
26
27
28	class ArrowStringMixin(ArrowColumnMixin):	2✔
29	# Reset for columns whose consumers need real items, e.g.
30	# LowCardinality dictionaries.
31	arrow_buffers_allowed = True	2✔
32
33	def _buffers_encoding_ok(self):	2✔
34	return self.encoding.lower() in ('utf-8', 'utf8')	2✔
35
36	def _use_arrow_buffers(self, buf):	2✔
37	return (	2✔
38	self.arrow_buffers_allowed and
39	self._buffers_encoding_ok() and
40	type(self.serialization) is CommonSerialization and
41	hasattr(buf, 'read_strings_arrow')
42	)
43
44	def _read_data(self, n_items, buf, nulls_map=None):	2✔
45	if self._use_arrow_buffers(buf):	2✔
46	offsets, data = buf.read_strings_arrow(n_items)	2✔
47	return ArrowStringBuffers(offsets, data, nulls_map=nulls_map)	2✔
48
49	return super(ArrowStringMixin, self)._read_data(	2✔
50	n_items, buf, nulls_map=nulls_map
51	)
52
53	def _wrap_items(self, items):	2✔
54	# Wrapping strings into an ndarray re-encodes them (unicode
55	# dtype). Arrow consumes the raw tuple directly; the ndarray
56	# is only needed for the nullable masked path.
57	if self.nullable:	2✔
58	return np.array(items, dtype=self.dtype)	2✔
59	return items	2✔
60
61
62	class ArrowStringColumn(ArrowStringMixin, NumpyStringColumn):	2✔
63	def read_items(self, n_items, buf):	2✔
64	return self._wrap_items(	2✔
65	buf.read_strings(n_items, encoding=self.encoding)
66	)
67
68
69	class ArrowByteStringColumn(ArrowStringMixin, NumpyByteStringColumn):	2✔
70	def _buffers_encoding_ok(self):	2✔
71	# No decoding happens for byte strings: buffers can be used
72	# regardless of encoding.
73	return True	2✔
74
75	def read_items(self, n_items, buf):	2✔
NEW 76	return self._wrap_items(buf.read_strings(n_items))	×
77
78
79	class ArrowFixedString(ArrowStringMixin, NumpyFixedString):	2✔
80	def _use_arrow_buffers(self, buf):	2✔
81	# Fixed strings are trimmed of zero bytes on read: buffers
82	# would keep the padding.
83	return False	2✔
84
85	def read_items(self, n_items, buf):	2✔
86	return self._wrap_items(buf.read_fixed_strings(	2✔
87	n_items, self.length, encoding=self.encoding
88	))
89
90
91	class ArrowByteFixedString(ArrowStringMixin, NumpyByteFixedString):	2✔
92	def _use_arrow_buffers(self, buf):	2✔
NEW 93	return False	×
94
95	def read_items(self, n_items, buf):	2✔
NEW 96	return self._wrap_items(	×
97	buf.read_fixed_strings(n_items, self.length)
98	)
99
100
101	def create_arrow_string_column(spec, column_options):	2✔
102	client_settings = column_options['context'].client_settings	2✔
103	strings_as_bytes = client_settings['strings_as_bytes']	2✔
104	encoding = client_settings.get(	2✔
105	'strings_encoding', NumpyStringColumn.default_encoding
106	)
107
108	if spec == 'String':	2✔
109	cls = ArrowByteStringColumn if strings_as_bytes \	2✔
110	else ArrowStringColumn
111	return cls(encoding=encoding, **column_options)	2✔
112	else:
113	length = int(spec[12:-1])	2✔
114	cls = ArrowByteFixedString if strings_as_bytes \	2✔
115	else ArrowFixedString
116	return cls(length, encoding=encoding, **column_options)	2✔

mymarilyn / clickhouse-driver / 29579092494

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous