ARROW-242: Support Timestamp Data Type

Change-Id: I8e5551cddec30888bbe37f1fc010774eba43f5dd
This commit is contained in:
Uwe L. Korn committed 2016-08-28 12:16:21 +02:00
1 parent 803afeb502
commit 34126b153d
14 files changed
+230 -21

No files matched your search

@@ -137,6 +137,15 @@ struct test_traits<Int64Type> {
const int64_t test_traits<Int64Type>::value(-1024);
template <>
struct test_traits<TimestampType> {
static constexpr ParquetType::type parquet_enum = ParquetType::INT64;
static constexpr LogicalType::type logical_enum = LogicalType::TIMESTAMP_MILLIS;
static int64_t const value;
};
const int64_t test_traits<TimestampType>::value(14695634030000);
template <>
struct test_traits<FloatType> {
static constexpr ParquetType::type parquet_enum = ParquetType::FLOAT;
@@ -248,7 +257,8 @@ class TestParquetIO : public ::testing::Test {
// Parquet version 1.0.
typedef ::testing::Types<BooleanType, UInt8Type, Int8Type, UInt16Type, Int16Type,
Int32Type, UInt64Type, Int64Type, FloatType, DoubleType, StringType> TestTypes;
Int32Type, UInt64Type, Int64Type, TimestampType, FloatType, DoubleType,
StringType> TestTypes;
TYPED_TEST_CASE(TestParquetIO, TestTypes);
+20 -3
View File
@@ -22,6 +22,7 @@
#include "arrow/test-util.h"
#include "arrow/type.h"
#include "arrow/types/datetime.h"
#include "arrow/types/decimal.h"
#include "arrow/util/status.h"
@@ -45,6 +46,9 @@ const auto INT64 = std::make_shared<Int64Type>();
const auto FLOAT = std::make_shared<FloatType>();
const auto DOUBLE = std::make_shared<DoubleType>();
const auto UTF8 = std::make_shared<StringType>();
const auto TIMESTAMP_MS = std::make_shared<TimestampType>(TimestampType::Unit::MILLI);
// TODO: This requires parquet-cpp implementing the MICROS enum value
// const auto TIMESTAMP_US = std::make_shared<TimestampType>(TimestampType::Unit::MICRO);
const auto BINARY = std::make_shared<ListType>(std::make_shared<Field>("", UINT8));
const auto DECIMAL_8_4 = std::make_shared<DecimalType>(8, 4);
@@ -89,6 +93,14 @@ TEST_F(TestConvertParquetSchema, ParquetFlatPrimitives) {
PrimitiveNode::Make("int64", Repetition::REQUIRED, ParquetType::INT64));
arrow_fields.push_back(std::make_shared<Field>("int64", INT64, false));
parquet_fields.push_back(PrimitiveNode::Make("timestamp", Repetition::REQUIRED,
ParquetType::INT64, LogicalType::TIMESTAMP_MILLIS));
arrow_fields.push_back(std::make_shared<Field>("timestamp", TIMESTAMP_MS, false));
// parquet_fields.push_back(PrimitiveNode::Make("timestamp", Repetition::REQUIRED,
// ParquetType::INT64, LogicalType::TIMESTAMP_MICROS));
// arrow_fields.push_back(std::make_shared<Field>("timestamp", TIMESTAMP_US, false));
parquet_fields.push_back(
PrimitiveNode::Make("float", Repetition::OPTIONAL, ParquetType::FLOAT));
arrow_fields.push_back(std::make_shared<Field>("float", FLOAT));
@@ -153,9 +165,6 @@ TEST_F(TestConvertParquetSchema, UnsupportedThings) {
unsupported_nodes.push_back(PrimitiveNode::Make(
"int32", Repetition::OPTIONAL, ParquetType::INT32, LogicalType::DATE));
unsupported_nodes.push_back(PrimitiveNode::Make(
"int64", Repetition::OPTIONAL, ParquetType::INT64, LogicalType::TIMESTAMP_MILLIS));
for (const NodePtr& node : unsupported_nodes) {
ASSERT_RAISES(NotImplemented, ConvertSchema({node}));
}
@@ -209,6 +218,14 @@ TEST_F(TestConvertArrowSchema, ParquetFlatPrimitives) {
PrimitiveNode::Make("int64", Repetition::REQUIRED, ParquetType::INT64));
arrow_fields.push_back(std::make_shared<Field>("int64", INT64, false));
parquet_fields.push_back(PrimitiveNode::Make("timestamp", Repetition::REQUIRED,
ParquetType::INT64, LogicalType::TIMESTAMP_MILLIS));
arrow_fields.push_back(std::make_shared<Field>("timestamp", TIMESTAMP_MS, false));
// parquet_fields.push_back(PrimitiveNode::Make("timestamp", Repetition::REQUIRED,
// ParquetType::INT64, LogicalType::TIMESTAMP_MICROS));
// arrow_fields.push_back(std::make_shared<Field>("timestamp", TIMESTAMP_US, false));
parquet_fields.push_back(
PrimitiveNode::Make("float", Repetition::OPTIONAL, ParquetType::FLOAT));
arrow_fields.push_back(std::make_shared<Field>("float", FLOAT));
+1
View File
@@ -368,6 +368,7 @@ Status FlatColumnReader::Impl::NextBatch(int batch_size, std::shared_ptr<Array>*
TYPED_BATCH_CASE(FLOAT, FloatType, ::parquet::FloatType)
TYPED_BATCH_CASE(DOUBLE, DoubleType, ::parquet::DoubleType)
TYPED_BATCH_CASE(STRING, StringType, ::parquet::ByteArrayType)
TYPED_BATCH_CASE(TIMESTAMP, TimestampType, ::parquet::Int64Type)
default:
return Status::NotImplemented(field_->type->ToString());
}
+11 -2
View File
@@ -52,6 +52,7 @@ const auto INT64 = std::make_shared<Int64Type>();
const auto FLOAT = std::make_shared<FloatType>();
const auto DOUBLE = std::make_shared<DoubleType>();
const auto UTF8 = std::make_shared<StringType>();
const auto TIMESTAMP_MS = std::make_shared<TimestampType>(TimestampType::Unit::MILLI);
const auto BINARY = std::make_shared<ListType>(std::make_shared<Field>("", UINT8));
TypePtr MakeDecimalType(const PrimitiveNode* node) {
@@ -133,6 +134,9 @@ static Status FromInt64(const PrimitiveNode* node, TypePtr* out) {
case LogicalType::DECIMAL:
*out = MakeDecimalType(node);
break;
case LogicalType::TIMESTAMP_MILLIS:
*out = TIMESTAMP_MS;
break;
default:
return Status::NotImplemented("Unhandled logical type for int64");
break;
@@ -289,10 +293,15 @@ Status FieldToNode(const std::shared_ptr<Field>& field,
type = ParquetType::INT32;
logical_type = LogicalType::DATE;
break;
case Type::TIMESTAMP:
case Type::TIMESTAMP: {
auto timestamp_type = static_cast<TimestampType*>(field->type.get());
if (timestamp_type->unit != TimestampType::Unit::MILLI) {
return Status::NotImplemented(
"Other timestamp units than millisecond are not yet support with parquet.");
}
type = ParquetType::INT64;
logical_type = LogicalType::TIMESTAMP_MILLIS;
break;
} break;
case Type::TIMESTAMP_DOUBLE:
type = ParquetType::INT64;
// This is specified as seconds since the UNIX epoch
+1
View File
@@ -240,6 +240,7 @@ Status FileWriter::Impl::WriteFlatColumnChunk(
TYPED_BATCH_CASE(INT32, Int32Type, ::parquet::Int32Type)
TYPED_BATCH_CASE(UINT64, UInt64Type, ::parquet::Int64Type)
TYPED_BATCH_CASE(INT64, Int64Type, ::parquet::Int64Type)
TYPED_BATCH_CASE(TIMESTAMP, TimestampType, ::parquet::Int64Type)
TYPED_BATCH_CASE(FLOAT, FloatType, ::parquet::FloatType)
TYPED_BATCH_CASE(DOUBLE, DoubleType, ::parquet::DoubleType)
default:
+2 -1
View File
@@ -51,6 +51,7 @@ Status MakeBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
BUILDER_CASE(INT32, Int32Builder);
BUILDER_CASE(UINT64, UInt64Builder);
BUILDER_CASE(INT64, Int64Builder);
BUILDER_CASE(TIMESTAMP, TimestampBuilder);
BUILDER_CASE(BOOL, BooleanBuilder);
@@ -105,7 +106,7 @@ Status MakePrimitiveArray(const TypePtr& type, int32_t length,
MAKE_PRIMITIVE_ARRAY_CASE(UINT64, UInt64Array);
MAKE_PRIMITIVE_ARRAY_CASE(INT64, Int64Array);
MAKE_PRIMITIVE_ARRAY_CASE(TIME, Int64Array);
MAKE_PRIMITIVE_ARRAY_CASE(TIMESTAMP, Int64Array);
MAKE_PRIMITIVE_ARRAY_CASE(TIMESTAMP, TimestampArray);
MAKE_PRIMITIVE_ARRAY_CASE(FLOAT, FloatArray);
MAKE_PRIMITIVE_ARRAY_CASE(DOUBLE, DoubleArray);
MAKE_PRIMITIVE_ARRAY_CASE(TIMESTAMP_DOUBLE, DoubleArray);
+9 -1
View File
@@ -34,15 +34,23 @@ struct DateType : public DataType {
static char const* name() { return "date"; }
};
struct TimestampType : public DataType {
struct ARROW_EXPORT TimestampType : public DataType {
enum class Unit : char { SECOND = 0, MILLI = 1, MICRO = 2, NANO = 3 };
typedef int64_t c_type;
static constexpr Type::type type_enum = Type::TIMESTAMP;
int value_size() const override { return sizeof(int64_t); }
Unit unit;
explicit TimestampType(Unit unit = Unit::MILLI)
: DataType(Type::TIMESTAMP), unit(unit) {}
TimestampType(const TimestampType& other) : TimestampType(other.unit) {}
virtual ~TimestampType() {}
std::string ToString() const override { return "timestamp"; }
static char const* name() { return "timestamp"; }
};
+1
View File
@@ -158,6 +158,7 @@ template class PrimitiveBuilder<Int8Type>;
template class PrimitiveBuilder<Int16Type>;
template class PrimitiveBuilder<Int32Type>;
template class PrimitiveBuilder<Int64Type>;
template class PrimitiveBuilder<TimestampType>;
template class PrimitiveBuilder<FloatType>;
template class PrimitiveBuilder<DoubleType>;
template class PrimitiveBuilder<BooleanType>;
+11
View File
@@ -26,6 +26,7 @@
#include "arrow/array.h"
#include "arrow/builder.h"
#include "arrow/type.h"
#include "arrow/types/datetime.h"
#include "arrow/util/bit-util.h"
#include "arrow/util/buffer.h"
#include "arrow/util/status.h"
@@ -100,6 +101,7 @@ NUMERIC_ARRAY_DECL(UInt32Array, UInt32Type);
NUMERIC_ARRAY_DECL(Int32Array, Int32Type);
NUMERIC_ARRAY_DECL(UInt64Array, UInt64Type);
NUMERIC_ARRAY_DECL(Int64Array, Int64Type);
NUMERIC_ARRAY_DECL(TimestampArray, TimestampType);
NUMERIC_ARRAY_DECL(FloatArray, FloatType);
NUMERIC_ARRAY_DECL(DoubleArray, DoubleType);
@@ -235,7 +237,15 @@ struct type_traits<Int64Type> {
static inline int bytes_required(int elements) { return elements * sizeof(int64_t); }
};
template <>
struct type_traits<TimestampType> {
typedef TimestampArray ArrayType;
static inline int bytes_required(int elements) { return elements * sizeof(int64_t); }
};
template <>
struct type_traits<FloatType> {
typedef FloatArray ArrayType;
@@ -260,6 +270,7 @@ typedef NumericBuilder<Int8Type> Int8Builder;
typedef NumericBuilder<Int16Type> Int16Builder;
typedef NumericBuilder<Int32Type> Int32Builder;
typedef NumericBuilder<Int64Type> Int64Builder;
typedef NumericBuilder<TimestampType> TimestampBuilder;
typedef NumericBuilder<FloatType> FloatBuilder;
typedef NumericBuilder<DoubleType> DoubleBuilder;
+37 -3
View File
@@ -19,6 +19,8 @@
# distutils: language = c++
# cython: embedsignature = True
import numpy as np
from pyarrow.includes.libarrow cimport *
cimport pyarrow.includes.pyarrow as pyarrow
@@ -186,6 +188,7 @@ cdef dict _array_classes = {
Type_DOUBLE: DoubleArray,
Type_LIST: ListArray,
Type_STRING: StringArray,
Type_TIMESTAMP: Int64Array,
}
cdef object box_arrow_array(const shared_ptr[CArray]& sp_array):
@@ -217,11 +220,28 @@ def from_pylist(object list_obj, DataType type=None):
return box_arrow_array(sp_array)
def from_pandas_series(object series, object mask=None):
def from_pandas_series(object series, object mask=None, timestamps_to_ms=False):
"""
Convert pandas.Series to an Arrow Array.
Parameters
----------
series: pandas.Series or numpy.ndarray
mask: pandas.Series or numpy.ndarray
array to mask null entries in the series
timestamps_to_ms: bool
Convert datetime columns to ms resolution. This is needed for
compability with other functionality like Parquet I/O which
only supports milliseconds.
"""
cdef:
shared_ptr[CArray] out
series_values = series_as_ndarray(series)
if series_values.dtype.type == np.datetime64 and timestamps_to_ms:
series_values = series_values.astype('datetime64[ms]')
if mask is None:
check_status(pyarrow.PandasToArrow(pyarrow.GetMemoryPool(),
@@ -234,14 +254,28 @@ def from_pandas_series(object series, object mask=None):
return box_arrow_array(out)
def from_pandas_dataframe(object df, name=None):
def from_pandas_dataframe(object df, name=None, timestamps_to_ms=False):
"""
Convert pandas.DataFrame to an Arrow Table
Parameters
----------
df: pandas.DataFrame
name: str
timestamps_to_ms: bool
Convert datetime columns to ms resolution. This is needed for
compability with other functionality like Parquet I/O which
only supports milliseconds.
"""
cdef:
list names = []
list arrays = []
for name in df.columns:
col = df[name]
arr = from_pandas_series(col)
arr = from_pandas_series(col, timestamps_to_ms=timestamps_to_ms)
names.append(name)
arrays.append(arr)
+1
View File
@@ -38,6 +38,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
Type_FLOAT" arrow::Type::FLOAT"
Type_DOUBLE" arrow::Type::DOUBLE"
Type_TIMESTAMP" arrow::Type::TIMESTAMP"
Type_STRING" arrow::Type::STRING"
Type_LIST" arrow::Type::LIST"
+22 -2
View File
@@ -33,8 +33,9 @@ class TestPandasConversion(unittest.TestCase):
def tearDown(self):
pass
def _check_pandas_roundtrip(self, df, expected=None):
table = A.from_pandas_dataframe(df)
def _check_pandas_roundtrip(self, df, expected=None,
timestamps_to_ms=False):
table = A.from_pandas_dataframe(df, timestamps_to_ms=timestamps_to_ms)
result = table.to_pandas()
if expected is None:
expected = df
@@ -164,6 +165,25 @@ class TestPandasConversion(unittest.TestCase):
expected = pd.DataFrame({'strings': values * repeats})
self._check_pandas_roundtrip(df, expected)
def test_timestamps_notimezone(self):
df = pd.DataFrame({
'datetime64': np.array([
'2007-07-13T01:23:34.123',
'2006-01-13T12:34:56.432',
'2010-08-13T05:46:57.437'],
dtype='datetime64[ms]')
})
self._check_pandas_roundtrip(df, timestamps_to_ms=True)
df = pd.DataFrame({
'datetime64': np.array([
'2007-07-13T01:23:34.123456789',
'2006-01-13T12:34:56.432539784',
'2010-08-13T05:46:57.437699912'],
dtype='datetime64[ns]')
})
self._check_pandas_roundtrip(df, timestamps_to_ms=False)
# def test_category(self):
# repeats = 1000
# values = [b'foo', None, u'bar', 'qux', np.nan]
+3 -1
View File
@@ -57,11 +57,13 @@ def test_pandas_parquet_2_0_rountrip(tmpdir):
'float32': np.arange(size, dtype=np.float32),
'float64': np.arange(size, dtype=np.float64),
'bool': np.random.randn(size) > 0,
# Pandas only support ns resolution, Arrow at the moment only ms
'datetime': np.arange("2016-01-01T00:00:00.001", size, dtype='datetime64[ms]'),
'str': [str(x) for x in range(size)],
'str_with_nulls': [None] + [str(x) for x in range(size - 2)] + [None]
})
filename = tmpdir.join('pandas_rountrip.parquet')
arrow_table = A.from_pandas_dataframe(df)
arrow_table = A.from_pandas_dataframe(df, timestamps_to_ms=True)
A.parquet.write_table(arrow_table, filename.strpath, version="2.0")
table_read = pyarrow.parquet.read_table(filename.strpath)
df_read = table_read.to_pandas()
+100 -7
View File
@@ -38,6 +38,7 @@ namespace pyarrow {
using arrow::Array;
using arrow::Column;
using arrow::DataType;
namespace util = arrow::util;
// ----------------------------------------------------------------------
@@ -50,7 +51,7 @@ struct npy_traits {
template <>
struct npy_traits<NPY_BOOL> {
typedef uint8_t value_type;
using ArrayType = arrow::BooleanArray;
using TypeClass = arrow::BooleanType;
static constexpr bool supports_nulls = false;
static inline bool isnull(uint8_t v) {
@@ -62,7 +63,7 @@ struct npy_traits<NPY_BOOL> {
template <> \
struct npy_traits<NPY_##TYPE> { \
typedef T value_type; \
using ArrayType = arrow::CapType##Array; \
using TypeClass = arrow::CapType##Type; \
\
static constexpr bool supports_nulls = false; \
static inline bool isnull(T v) { \
@@ -82,7 +83,7 @@ NPY_INT_DECL(UINT64, UInt64, uint64_t);
template <>
struct npy_traits<NPY_FLOAT32> {
typedef float value_type;
using ArrayType = arrow::FloatArray;
using TypeClass = arrow::FloatType;
static constexpr bool supports_nulls = true;
@@ -94,7 +95,7 @@ struct npy_traits<NPY_FLOAT32> {
template <>
struct npy_traits<NPY_FLOAT64> {
typedef double value_type;
using ArrayType = arrow::DoubleArray;
using TypeClass = arrow::DoubleType;
static constexpr bool supports_nulls = true;
@@ -103,6 +104,22 @@ struct npy_traits<NPY_FLOAT64> {
}
};
template <>
struct npy_traits<NPY_DATETIME> {
typedef double value_type;
using TypeClass = arrow::TimestampType;
static constexpr bool supports_nulls = true;
static inline bool isnull(int64_t v) {
// NaT = -2**63
// = -0x8000000000000000
// = -9223372036854775808;
// = std::numeric_limits<int64_t>::min()
return v == std::numeric_limits<int64_t>::min();
}
};
template <>
struct npy_traits<NPY_OBJECT> {
typedef PyObject* value_type;
@@ -206,6 +223,8 @@ class ArrowSerializer {
return Status::OK();
}
Status MakeDataType(std::shared_ptr<DataType>* out);
arrow::MemoryPool* pool_;
PyArrayObject* arr_;
@@ -253,6 +272,39 @@ static int64_t ValuesToBitmap(const void* data, int64_t length, uint8_t* bitmap)
return null_count;
}
template <int TYPE>
inline Status ArrowSerializer<TYPE>::MakeDataType(std::shared_ptr<DataType>* out) {
out->reset(new typename npy_traits<TYPE>::TypeClass());
return Status::OK();
}
template <>
inline Status ArrowSerializer<NPY_DATETIME>::MakeDataType(std::shared_ptr<DataType>* out) {
PyArray_Descr* descr = PyArray_DESCR(arr_);
auto date_dtype = reinterpret_cast<PyArray_DatetimeDTypeMetaData*>(descr->c_metadata);
arrow::TimestampType::Unit unit;
switch (date_dtype->meta.base) {
case NPY_FR_s:
unit = arrow::TimestampType::Unit::SECOND;
break;
case NPY_FR_ms:
unit = arrow::TimestampType::Unit::MILLI;
break;
case NPY_FR_us:
unit = arrow::TimestampType::Unit::MICRO;
break;
case NPY_FR_ns:
unit = arrow::TimestampType::Unit::NANO;
break;
default:
return Status::ValueError("Unknown NumPy datetime unit");
}
out->reset(new arrow::TimestampType(unit));
return Status::OK();
}
template <int TYPE>
inline Status ArrowSerializer<TYPE>::Convert(std::shared_ptr<Array>* out) {
typedef npy_traits<TYPE> traits;
@@ -269,9 +321,9 @@ inline Status ArrowSerializer<TYPE>::Convert(std::shared_ptr<Array>* out) {
}
RETURN_NOT_OK(ConvertData());
*out = std::make_shared<typename traits::ArrayType>(length_, data_, null_count,
null_bitmap_);
std::shared_ptr<DataType> type;
RETURN_NOT_OK(MakeDataType(&type));
RETURN_ARROW_NOT_OK(MakePrimitiveArray(type, length_, data_, null_count, null_bitmap_, out));
return Status::OK();
}
@@ -402,6 +454,7 @@ Status PandasMaskedToArrow(arrow::MemoryPool* pool, PyObject* ao, PyObject* mo,
TO_ARROW_CASE(UINT64);
TO_ARROW_CASE(FLOAT32);
TO_ARROW_CASE(FLOAT64);
TO_ARROW_CASE(DATETIME);
TO_ARROW_CASE(OBJECT);
default:
std::stringstream ss;
@@ -476,6 +529,17 @@ struct arrow_traits<arrow::Type::DOUBLE> {
typedef typename npy_traits<NPY_FLOAT64>::value_type T;
};
template <>
struct arrow_traits<arrow::Type::TIMESTAMP> {
static constexpr int npy_type = NPY_DATETIME;
static constexpr bool supports_nulls = true;
static constexpr int64_t na_value = std::numeric_limits<int64_t>::min();
static constexpr bool is_boolean = false;
static constexpr bool is_integer = true;
static constexpr bool is_floating = false;
typedef typename npy_traits<NPY_DATETIME>::value_type T;
};
template <>
struct arrow_traits<arrow::Type::STRING> {
static constexpr int npy_type = NPY_OBJECT;
@@ -494,6 +558,30 @@ static inline PyObject* make_pystring(const uint8_t* data, int32_t length) {
#endif
}
inline void set_numpy_metadata(int type, DataType* datatype, PyArrayObject* out) {
if (type == NPY_DATETIME) {
auto timestamp_type = static_cast<arrow::TimestampType*>(datatype);
// We only support ms resolution at the moment
PyArray_Descr* descr = PyArray_DESCR(out);
auto date_dtype = reinterpret_cast<PyArray_DatetimeDTypeMetaData*>(descr->c_metadata);
switch (timestamp_type->unit) {
case arrow::TimestampType::Unit::SECOND:
date_dtype->meta.base = NPY_FR_s;
break;
case arrow::TimestampType::Unit::MILLI:
date_dtype->meta.base = NPY_FR_ms;
break;
case arrow::TimestampType::Unit::MICRO:
date_dtype->meta.base = NPY_FR_us;
break;
case arrow::TimestampType::Unit::NANO:
date_dtype->meta.base = NPY_FR_ns;
break;
}
}
}
template <int TYPE>
class ArrowDeserializer {
public:
@@ -522,6 +610,8 @@ class ArrowDeserializer {
return Status::OK();
}
set_numpy_metadata(type, col_->type().get(), out_);
return Status::OK();
}
@@ -538,6 +628,8 @@ class ArrowDeserializer {
return Status::OK();
}
set_numpy_metadata(type, col_->type().get(), out_);
if (PyArray_SetBaseObject(out_, py_ref_) == -1) {
// Error occurred, trust that SetBaseObject set the error state
return Status::OK();
@@ -713,6 +805,7 @@ Status ArrowToPandas(const std::shared_ptr<Column>& col, PyObject* py_ref,
FROM_ARROW_CASE(FLOAT);
FROM_ARROW_CASE(DOUBLE);
FROM_ARROW_CASE(STRING);
FROM_ARROW_CASE(TIMESTAMP);
default:
return Status::NotImplemented("Arrow type reading not implemented");
}