Skip to content

Commit b694825

Browse files
pitroufsaintjacques
pitrou
authored and
fsaintjacques
committed
ARROW-8711: [Python] Expose timestamp_parsers in csv.ConvertOptions
Closes apache#7223 from pitrou/ARROW-8711-py-csv-timestamp-parsing Authored-by: Antoine Pitrou <antoine@python.org> Signed-off-by: François Saint-Jacques <fsaintjacques@gmail.com>
1 parent 53e8438 commit b694825

10 files changed

Lines changed: 186 additions & 33 deletions

File tree

‎cpp/src/arrow/csv/options.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ struct ARROW_EXPORT ConvertOptions {
105105
/// arrow/util/value_parsing.h. More than one parser can be specified, and
106106
/// the CSV conversion logic will try parsing values starting from the
107107
/// beginning of this vector. If no parsers are specified, we use the default
108-
/// built-in ISO-8601 parser
108+
/// built-in ISO-8601 parser.
109109
std::vector<std::shared_ptr<TimestampParser>> timestamp_parsers;
110110

111111
/// Create conversion options with default values, including conventional

‎cpp/src/arrow/util/value_parsing.cc‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,8 @@ bool StringToFloat(const char* s, size_t length, double* out) {
8989
// ----------------------------------------------------------------------
9090
// strptime-like parsing
9191

92+
namespace {
93+
9294
class StrptimeTimestampParser : public TimestampParser {
9395
public:
9496
explicit StrptimeTimestampParser(std::string format) : format_(std::move(format)) {}
@@ -100,6 +102,10 @@ class StrptimeTimestampParser : public TimestampParser {
100102
/*allow_trailing_chars=*/false, out_unit, out);
101103
}
102104

105+
const char* kind() const override { return "strptime"; }
106+
107+
const char* format() const override { return format_.c_str(); }
108+
103109
private:
104110
std::string format_;
105111
};
@@ -112,10 +118,15 @@ class ISO8601Parser : public TimestampParser {
112118
int64_t* out) const override {
113119
return ParseTimestampISO8601(s, length, out_unit, out);
114120
}
121+
122+
const char* kind() const override { return "iso8601"; }
115123
};
116124

125+
} // namespace
117126
} // namespace internal
118127

128+
const char* TimestampParser::format() const { return ""; }
129+
119130
std::shared_ptr<TimestampParser> TimestampParser::MakeStrptime(std::string format) {
120131
return std::make_shared<internal::StrptimeTimestampParser>(std::move(format));
121132
}

‎cpp/src/arrow/util/value_parsing.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@ class ARROW_EXPORT TimestampParser {
4545
virtual bool operator()(const char* s, size_t length, TimeUnit::type out_unit,
4646
int64_t* out) const = 0;
4747

48+
virtual const char* kind() const = 0;
49+
50+
virtual const char* format() const;
51+
4852
/// \brief Create a TimestampParser that recognizes strptime-like format strings
4953
static std::shared_ptr<TimestampParser> MakeStrptime(std::string format);
5054

‎python/pyarrow/_csv.pyx‎

Lines changed: 75 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
# cython: embedsignature = True
2121
# cython: language_level = 3
2222

23+
from cython.operator cimport dereference as deref
2324

2425
from pyarrow.includes.common cimport *
2526
from pyarrow.includes.libarrow cimport *
@@ -69,7 +70,7 @@ cdef class ReadOptions:
6970
# Avoid mistakingly creating attributes
7071
__slots__ = ()
7172

72-
def __init__(self, use_threads=None, block_size=None, skip_rows=None,
73+
def __init__(self, *, use_threads=None, block_size=None, skip_rows=None,
7374
column_names=None, autogenerate_column_names=None):
7475
self.options = CCSVReadOptions.Defaults()
7576
if use_threads is not None:
@@ -176,7 +177,7 @@ cdef class ParseOptions:
176177
"""
177178
__slots__ = ()
178179

179-
def __init__(self, delimiter=None, quote_char=None, double_quote=None,
180+
def __init__(self, *, delimiter=None, quote_char=None, double_quote=None,
180181
escape_char=None, newlines_in_values=None,
181182
ignore_empty_lines=None):
182183
self.options = CCSVParseOptions.Defaults()
@@ -296,10 +297,31 @@ cdef class ParseOptions:
296297
out.options = options
297298
return out
298299

299-
def __reduce__(self):
300-
return ParseOptions, (self.delimiter, self.quote_char,
301-
self.double_quote, self.escape_char,
302-
self.newlines_in_values, self.ignore_empty_lines)
300+
def __getstate__(self):
301+
return (self.delimiter, self.quote_char, self.double_quote,
302+
self.escape_char, self.newlines_in_values,
303+
self.ignore_empty_lines)
304+
305+
def __setstate__(self, state):
306+
(self.delimiter, self.quote_char, self.double_quote,
307+
self.escape_char, self.newlines_in_values,
308+
self.ignore_empty_lines) = state
309+
310+
311+
cdef class _ISO8601:
312+
"""
313+
A special object indicating ISO-8601 parsing.
314+
"""
315+
__slots__ = ()
316+
317+
def __str__(self):
318+
return 'ISO8601'
319+
320+
def __eq__(self, other):
321+
return isinstance(other, _ISO8601)
322+
323+
324+
ISO8601 = _ISO8601()
303325

304326

305327
cdef class ConvertOptions:
@@ -324,6 +346,11 @@ cdef class ConvertOptions:
324346
false_values: list, optional
325347
A sequence of strings that denote false booleans in the data
326348
(defaults are appropriate in most cases).
349+
timestamp_parsers: list, optional
350+
A sequence of strptime()-compatible format strings, tried in order
351+
when attempting to infer or convert timestamp values (the special
352+
value ISO8601() can also be given). By default, a fast built-in
353+
ISO-8601 parser is used.
327354
strings_can_be_null: bool, optional (default False)
328355
Whether string / binary columns can have null values.
329356
If true, then strings in null_values are considered null for
@@ -359,11 +386,11 @@ cdef class ConvertOptions:
359386
# Avoid mistakingly creating attributes
360387
__slots__ = ()
361388

362-
def __init__(self, check_utf8=None, column_types=None, null_values=None,
389+
def __init__(self, *, check_utf8=None, column_types=None, null_values=None,
363390
true_values=None, false_values=None,
364391
strings_can_be_null=None, include_columns=None,
365392
include_missing_columns=None, auto_dict_encode=None,
366-
auto_dict_max_cardinality=None):
393+
auto_dict_max_cardinality=None, timestamp_parsers=None):
367394
self.options = CCSVConvertOptions.Defaults()
368395
if check_utf8 is not None:
369396
self.check_utf8 = check_utf8
@@ -385,6 +412,8 @@ cdef class ConvertOptions:
385412
self.auto_dict_encode = auto_dict_encode
386413
if auto_dict_max_cardinality is not None:
387414
self.auto_dict_max_cardinality = auto_dict_max_cardinality
415+
if timestamp_parsers is not None:
416+
self.timestamp_parsers = timestamp_parsers
388417

389418
@property
390419
def check_utf8(self):
@@ -526,6 +555,44 @@ cdef class ConvertOptions:
526555
def include_missing_columns(self, value):
527556
self.options.include_missing_columns = value
528557

558+
@property
559+
def timestamp_parsers(self):
560+
"""
561+
A sequence of strptime()-compatible format strings, tried in order
562+
when attempting to infer or convert timestamp values (the special
563+
value ISO8601() can also be given). By default, a fast built-in
564+
ISO-8601 parser is used.
565+
"""
566+
cdef:
567+
shared_ptr[CTimestampParser] c_parser
568+
c_string kind
569+
570+
parsers = []
571+
for c_parser in self.options.timestamp_parsers:
572+
kind = deref(c_parser).kind()
573+
if kind == b'strptime':
574+
parsers.append(frombytes(deref(c_parser).format()))
575+
else:
576+
assert kind == b'iso8601'
577+
parsers.append(ISO8601)
578+
579+
return parsers
580+
581+
@timestamp_parsers.setter
582+
def timestamp_parsers(self, value):
583+
cdef:
584+
vector[shared_ptr[CTimestampParser]] c_parsers
585+
586+
for v in value:
587+
if isinstance(v, str):
588+
c_parsers.push_back(CTimestampParser.MakeStrptime(tobytes(v)))
589+
elif v == ISO8601:
590+
c_parsers.push_back(CTimestampParser.MakeISO8601())
591+
else:
592+
raise TypeError("Expected list of str or ISO8601 objects")
593+
594+
self.options.timestamp_parsers = move(c_parsers)
595+
529596

530597
cdef _get_reader(input_file, shared_ptr[CInputStream]* out):
531598
use_memory_map = False

‎python/pyarrow/array.pxi‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1565,7 +1565,7 @@ cdef class UnionArray(Array):
15651565
For dense unions, the returned array is unchanged.
15661566
"""
15671567
cdef shared_ptr[CArray] result
1568-
result = (<CUnionArray*> self.ap).child(pos)
1568+
result = (<CUnionArray*> self.ap).field(pos)
15691569
if result != NULL:
15701570
return pyarrow_wrap_array(result)
15711571
raise KeyError("UnionArray does not have child {}".format(pos))

‎python/pyarrow/csv.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,5 +17,5 @@
1717

1818

1919
from pyarrow._csv import ( # noqa
20-
ReadOptions, ParseOptions, ConvertOptions, open_csv, read_csv,
21-
CSVStreamingReader)
20+
ReadOptions, ParseOptions, ConvertOptions, ISO8601,
21+
open_csv, read_csv, CSVStreamingReader)

‎python/pyarrow/includes/libarrow.pxd‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -111,11 +111,11 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
111111

112112
c_bool Equals(const CDataType& other)
113113

114-
shared_ptr[CField] child(int i)
114+
shared_ptr[CField] field(int i)
115115

116-
const vector[shared_ptr[CField]] children()
116+
const vector[shared_ptr[CField]] fields()
117117

118-
int num_children()
118+
int num_fields()
119119

120120
CDataTypeLayout layout()
121121

@@ -572,7 +572,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
572572
int32_t value_offset(int i)
573573
shared_ptr[CBuffer] value_offsets()
574574
int child_id(int64_t index)
575-
shared_ptr[CArray] child(int pos)
575+
shared_ptr[CArray] field(int pos)
576576
const CArray* UnsafeChild(int pos)
577577
UnionMode mode()
578578

@@ -1362,6 +1362,18 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil:
13621362
CStatus Read(const vector[c_string] names, shared_ptr[CTable]* out)
13631363

13641364

1365+
cdef extern from 'arrow/util/value_parsing.h' namespace 'arrow' nogil:
1366+
cdef cppclass CTimestampParser" arrow::TimestampParser":
1367+
const char* kind() const
1368+
const char* format() const
1369+
1370+
@staticmethod
1371+
shared_ptr[CTimestampParser] MakeStrptime(c_string format)
1372+
1373+
@staticmethod
1374+
shared_ptr[CTimestampParser] MakeISO8601()
1375+
1376+
13651377
cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil:
13661378

13671379
cdef cppclass CCSVParseOptions" arrow::csv::ParseOptions":
@@ -1384,6 +1396,7 @@ cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil:
13841396
vector[c_string] true_values
13851397
vector[c_string] false_values
13861398
c_bool strings_can_be_null
1399+
vector[shared_ptr[CTimestampParser]] timestamp_parsers
13871400

13881401
c_bool auto_dict_encode
13891402
int32_t auto_dict_max_cardinality

‎python/pyarrow/scalar.pxi‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -849,7 +849,7 @@ cdef class UnionValue(ArrayValue):
849849

850850
cdef getitem(self, int64_t i):
851851
cdef int child_id = self.ap.child_id(i)
852-
cdef shared_ptr[CArray] child = self.ap.child(child_id)
852+
cdef shared_ptr[CArray] child = self.ap.field(child_id)
853853
if self.ap.mode() == _UnionMode_SPARSE:
854854
return box_scalar(self.type[child_id].type, child, i)
855855
else:
@@ -916,16 +916,15 @@ cdef class StructValue(ArrayValue):
916916
Return this value as a Python dict.
917917
"""
918918
cdef:
919-
vector[shared_ptr[CField]] child_fields = self.type.type.children()
919+
vector[shared_ptr[CField]] child_fields = self.type.type.fields()
920920

921921
wrapped_arrays = [pyarrow_wrap_array(self.ap.field(i))
922922
for i in range(self.ap.num_fields())]
923923
child_names = [child.get().name() for child in child_fields]
924924
# Return the struct as a dict
925925
return {
926926
frombytes(name): child_array[self.index].as_py()
927-
for name, child_array in
928-
zip(child_names, wrapped_arrays)
927+
for name, child_array in zip(child_names, wrapped_arrays)
929928
}
930929

931930

0 commit comments

Comments
 (0)