From 7e8b0297af716c87a31dd954f07485695b15f956 Mon Sep 17 00:00:00 2001 From: Edmond Abraham <54119952+gitedmond@users.noreply.github.com> Date: Sat, 10 Oct 2026 22:12:14 -0400 Subject: [PATCH] GH-50142: [C++][Python] Honor Parquet time coercion and format version --- .../parquet/arrow/arrow_reader_writer_test.cc | 35 +++++- cpp/src/parquet/arrow/arrow_schema_test.cc | 32 +++++ cpp/src/parquet/arrow/schema.cc | 63 +++++++--- cpp/src/parquet/column_writer.cc | 119 +++++++++++++++--- cpp/src/parquet/properties.h | 10 +- .../python/parquet/parquet_type_handling.rst | 20 +-- python/pyarrow/parquet/core.py | 12 +- python/pyarrow/tests/parquet/test_datetime.py | 110 ++++++++++++++++ 8 files changed, 352 insertions(+), 49 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index d858962d4728..9ded03867bf5 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -1873,7 +1873,8 @@ void MakeDateTimeTypesTable(std::shared_ptr* out, bool expected_micro = f auto f3_x = field("f3", ::arrow::timestamp(TimeUnit::MICRO)); auto f4 = field("f4", ::arrow::time32(TimeUnit::MILLI)); auto f5 = field("f5", ::arrow::time64(TimeUnit::MICRO)); - auto f6 = field("f6", ::arrow::time64(TimeUnit::NANO)); + auto f6 = + field("f6", ::arrow::time64(expected_micro ? TimeUnit::MICRO : TimeUnit::NANO)); std::shared_ptr<::arrow::Schema> schema( new ::arrow::Schema({f0, f1, f2, (expected_micro ? f3_x : f3), f4, f5, f6})); @@ -1899,7 +1900,8 @@ void MakeDateTimeTypesTable(std::shared_ptr
* out, bool expected_micro = f &a3_x); ArrayFromVector<::arrow::Time32Type, int32_t>(f4->type(), is_valid, t32_values, &a4); ArrayFromVector<::arrow::Time64Type, int64_t>(f5->type(), is_valid, t64_us_values, &a5); - ArrayFromVector<::arrow::Time64Type, int64_t>(f6->type(), is_valid, t64_ns_values, &a6); + ArrayFromVector<::arrow::Time64Type, int64_t>( + f6->type(), is_valid, expected_micro ? t64_us_values : t64_ns_values, &a6); *out = Table::Make(schema, {a0, a1, a2, expected_micro ? a3_x : a3, a4, a5, a6}); } @@ -1927,6 +1929,35 @@ TEST(TestArrowReadWrite, DateTimeTypes) { CheckConfiguredRoundtrip(table, expected, parquet_version_2_4_properties)); } +TEST(TestArrowReadWrite, WriteUnannotatedTimeArrays) { + for (bool wide : {false, true}) { + ASSERT_OK_AND_ASSIGN(auto sink, ::arrow::io::BufferOutputStream::Create()); + auto node = PrimitiveNode::Make("time", Repetition::OPTIONAL, + wide ? ParquetType::INT64 : ParquetType::INT32); + auto schema = std::static_pointer_cast( + GroupNode::Make("schema", Repetition::REQUIRED, {node})); + auto array = ArrayFromJSON( + wide ? ::arrow::time64(TimeUnit::NANO) : ::arrow::time32(TimeUnit::SECOND), + "[1000]"); + auto properties = default_arrow_writer_properties(); + ArrowWriteContext ctx(::arrow::default_memory_pool(), properties.get()); + auto writer = ParquetFileWriter::Open(sink, schema); + auto row_group = writer->AppendRowGroup(); + const int16_t def_level = 1; + ASSERT_OK( + row_group->NextColumn()->WriteArrow(&def_level, nullptr, 1, *array, &ctx, true)); + row_group->Close(); + writer->Close(); + ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + ASSERT_OK_AND_ASSIGN(auto reader, OpenFile(std::make_shared(buffer), + ::arrow::default_memory_pool())); + ASSERT_OK_AND_ASSIGN(auto result, reader->ReadTable()); + auto expected = ArrayFromJSON(wide ? ::arrow::int64() : ::arrow::int32(), + wide ? "[1000]" : "[1000000]"); + ::arrow::AssertArraysEqual(*expected, *result->column(0)->chunk(0)); + } +} + TEST(TestArrowReadWrite, UseDeprecatedInt96) { using ::arrow::ArrayFromVector; using ::arrow::field; diff --git a/cpp/src/parquet/arrow/arrow_schema_test.cc b/cpp/src/parquet/arrow/arrow_schema_test.cc index 0ec943841616..85e33152acc8 100644 --- a/cpp/src/parquet/arrow/arrow_schema_test.cc +++ b/cpp/src/parquet/arrow/arrow_schema_test.cc @@ -1871,6 +1871,38 @@ TEST_F(TestConvertArrowSchema, ParquetFlatDecimals) { ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); } +TEST_F(TestConvertArrowSchema, ParquetTimeCoercion) { + auto schema = + ::arrow::schema({::arrow::field("seconds", ::arrow::time32(TimeUnit::SECOND)), + ::arrow::field("nanos", ::arrow::time64(TimeUnit::NANO))}); + auto arrow_properties = ArrowWriterProperties::Builder() + .coerce_timestamps(TimeUnit::NANO) + ->set_time_adjusted_to_utc(true) + ->build(); + for (auto version : {ParquetVersion::PARQUET_1_0, ParquetVersion::PARQUET_2_4, + ParquetVersion::PARQUET_2_6}) { + auto properties = WriterProperties::Builder().version(version)->build(); + auto status = + ToParquetSchema(schema.get(), *properties, *arrow_properties, &result_schema_); + if (version != ParquetVersion::PARQUET_2_6) { + ASSERT_RAISES(NotImplemented, status); + continue; + } + ASSERT_OK(status); + const auto logical_type = LogicalType::Time(true, LogicalType::TimeUnit::NANOS); + for (int i = 0; i < schema->num_fields(); ++i) { + const auto* column = result_schema_->Column(i); + ASSERT_EQ(column->physical_type(), ParquetType::INT64); + ASSERT_TRUE(column->logical_type()->Equals(*logical_type)); + } + } + arrow_properties = + ArrowWriterProperties::Builder().coerce_timestamps(TimeUnit::SECOND)->build(); + ASSERT_RAISES(NotImplemented, + ToParquetSchema(schema.get(), *default_writer_properties(), + *arrow_properties, &result_schema_)); +} + TEST_F(TestConvertArrowSchema, ParquetTimeAdjustedToUTC) { // Verify Parquet Time types have the appropriate isAdjustedToUTC value, depending // on the return value of ArrowWriterProperties::write_time_adjusted_to_utc() diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc index bc4de6c39b5e..cac4d423465a 100644 --- a/cpp/src/parquet/arrow/schema.cc +++ b/cpp/src/parquet/arrow/schema.cc @@ -204,6 +204,48 @@ static std::shared_ptr TimestampLogicalTypeFromArrowTimestamp return LogicalType::None(); } +static Status GetTimeMetadata(const ::arrow::TimeType& type, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, + ParquetType::type* physical_type, + std::shared_ptr* logical_type) { + const bool coerce = arrow_properties.coerce_timestamps_enabled(); + auto unit = coerce ? arrow_properties.coerce_timestamps_unit() : type.unit(); + const bool supports_nanos = properties.version() == ParquetVersion::PARQUET_2_6; + if (!coerce) { + if (unit == ::arrow::TimeUnit::SECOND) { + unit = ::arrow::TimeUnit::MILLI; + } else if (unit == ::arrow::TimeUnit::NANO && !supports_nanos) { + unit = ::arrow::TimeUnit::MICRO; + } + } + + *physical_type = + unit == ::arrow::TimeUnit::MILLI ? ParquetType::INT32 : ParquetType::INT64; + LogicalType::TimeUnit::unit parquet_unit; + switch (unit) { + case ::arrow::TimeUnit::MILLI: + parquet_unit = LogicalType::TimeUnit::MILLIS; + break; + case ::arrow::TimeUnit::MICRO: + parquet_unit = LogicalType::TimeUnit::MICROS; + break; + case ::arrow::TimeUnit::NANO: + if (!supports_nanos) { + return Status::NotImplemented("Nanosecond times require Parquet version 2.6"); + } + parquet_unit = LogicalType::TimeUnit::NANOS; + break; + case ::arrow::TimeUnit::SECOND: + return Status::NotImplemented("Parquet does not support times in seconds"); + default: + return Status::NotImplemented("Unsupported Arrow time unit"); + } + *logical_type = + LogicalType::Time(arrow_properties.write_time_adjusted_to_utc(), parquet_unit); + return Status::OK(); +} + static Status GetTimestampMetadata(const ::arrow::TimestampType& type, const WriterProperties& properties, const ArrowWriterProperties& arrow_properties, @@ -423,24 +465,11 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, properties, arrow_properties, &type, &logical_type)); break; case ArrowTypeId::TIME32: - type = ParquetType::INT32; - logical_type = LogicalType::Time( - /*is_adjusted_to_utc=*/arrow_properties.write_time_adjusted_to_utc(), - LogicalType::TimeUnit::MILLIS); + case ArrowTypeId::TIME64: + RETURN_NOT_OK( + GetTimeMetadata(checked_cast(*field->type()), + properties, arrow_properties, &type, &logical_type)); break; - case ArrowTypeId::TIME64: { - type = ParquetType::INT64; - auto time_type = static_cast<::arrow::Time64Type*>(field->type().get()); - if (time_type->unit() == ::arrow::TimeUnit::NANO) { - logical_type = LogicalType::Time( - /*is_adjusted_to_utc=*/arrow_properties.write_time_adjusted_to_utc(), - LogicalType::TimeUnit::NANOS); - } else { - logical_type = LogicalType::Time( - /*is_adjusted_to_utc=*/arrow_properties.write_time_adjusted_to_utc(), - LogicalType::TimeUnit::MICROS); - } - } break; case ArrowTypeId::DURATION: type = ParquetType::INT64; break; diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index 3296af62f0c4..45f129da3db5 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -45,6 +46,7 @@ #include "arrow/util/key_value_metadata.h" #include "arrow/util/logging_internal.h" #include "arrow/util/rle_encoding_internal.h" +#include "arrow/util/time.h" #include "arrow/util/type_traits.h" #include "arrow/util/unreachable.h" #include "arrow/visit_array_inline.h" @@ -1546,6 +1548,11 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, return Status::OK(); } + template + Status WriteArrowTimes(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_levels, const ::arrow::Array& array, + ArrowWriteContext* ctx, bool maybe_parent_nulls); + Status WriteArrowTimestamps(const int16_t* def_levels, const int16_t* rep_levels, int64_t num_levels, const ::arrow::Array& values, ArrowWriteContext* ctx, bool maybe_parent_nulls) { @@ -2230,21 +2237,95 @@ struct SerializeFunctor< } }; -template <> -struct SerializeFunctor { - Status Serialize(const ::arrow::Time32Array& array, ArrowWriteContext*, int32_t* out) { - const int32_t* input = array.raw_values(); - const auto& type = static_cast(*array.type()); - if (type.unit() == ::arrow::TimeUnit::SECOND) { - for (int64_t i = 0; i < array.length(); i++) { - out[i] = input[i] * 1000; +template +template +Status TypedColumnWriterImpl::WriteArrowTimes( + const int16_t* def_levels, const int16_t* rep_levels, int64_t num_levels, + const ::arrow::Array& array, ArrowWriteContext* ctx, bool maybe_parent_nulls) { + using ArrayType = typename ::arrow::TypeTraits::ArrayType; + const auto& values = checked_cast(array); + const auto& source_type = checked_cast(*array.type()); + const auto& logical_type = this->descr()->logical_type(); + auto target_unit = source_type.unit() == ::arrow::TimeUnit::SECOND + ? ::arrow::TimeUnit::MILLI + : source_type.unit(); + if (logical_type->is_time()) { + switch (checked_cast(*logical_type).time_unit()) { + case LogicalType::TimeUnit::MILLIS: + target_unit = ::arrow::TimeUnit::MILLI; + break; + case LogicalType::TimeUnit::MICROS: + target_unit = ::arrow::TimeUnit::MICRO; + break; + case LogicalType::TimeUnit::NANOS: + target_unit = ::arrow::TimeUnit::NANO; + break; + default: + return Status::Invalid("Unsupported Parquet time unit"); + } + } else if constexpr (sizeof(T) != sizeof(typename ArrowType::c_type)) { + ARROW_UNSUPPORTED(); + } + if (source_type.unit() == target_unit) { + return WriteArrowZeroCopy(def_levels, rep_levels, num_levels, array, ctx, + maybe_parent_nulls); + } + + const auto target_type = target_unit == ::arrow::TimeUnit::MILLI + ? ::arrow::time32(target_unit) + : ::arrow::time64(target_unit); + const auto [operation, factor] = + ::arrow::util::GetTimestampConversion(source_type.unit(), target_unit); + const bool allow_truncation = ctx->properties->coerce_timestamps_enabled() && + ctx->properties->truncated_timestamps_allowed(); + T* out = nullptr; + RETURN_NOT_OK(ctx->GetScratchData(array.length(), &out)); + int64_t level_index = 0; + for (int64_t i = 0; i < array.length(); ++i) { + out[i] = 0; + bool valid = array.IsValid(i); + if (maybe_parent_nulls) { + // Definition levels also mask children beneath null parents. + while (level_index < num_levels && + def_levels[level_index] < level_info_.repeated_ancestor_def_level) { + ++level_index; } - } else { - std::copy(input, input + array.length(), out); + valid = + level_index < num_levels && def_levels[level_index++] == level_info_.def_level; + } + if (!valid) { + continue; + } + const int64_t value = values.Value(i); + int64_t converted; + if (operation == ::arrow::util::DIVIDE) { + if (!allow_truncation && value % factor != 0) { + return Status::Invalid("Casting from ", source_type.ToString(), " to ", + target_type->ToString(), " would lose data: ", value); + } + converted = value / factor; + } else if (::arrow::internal::MultiplyWithOverflowGeneric(value, factor, + &converted)) { + return Status::Invalid("Integer overflow when casting time value ", value); } - return Status::OK(); + if (converted < std::numeric_limits::min() || + converted > std::numeric_limits::max()) { + return Status::Invalid("Integer overflow when casting time value ", value); + } + out[i] = static_cast(converted); } -}; + + const bool no_nulls = + this->descr()->schema_node()->is_required() || array.null_count() == 0; + if (!maybe_parent_nulls && no_nulls) { + PARQUET_CATCH_NOT_OK(WriteBatchInternal(num_levels, def_levels, rep_levels, out)); + } else { + PARQUET_CATCH_NOT_OK(WriteBatchSpacedInternal(num_levels, def_levels, rep_levels, + array.null_bitmap_data(), + array.offset(), out)); + } + return Status::OK(); +} template <> Status TypedColumnWriterImpl::WriteArrowDense( @@ -2263,7 +2344,12 @@ Status TypedColumnWriterImpl::WriteArrowDense( WRITE_ZERO_COPY_CASE(INT32) WRITE_ZERO_COPY_CASE(DATE32) WRITE_SERIALIZE_CASE(DATE64) - WRITE_SERIALIZE_CASE(TIME32) + case ::arrow::Type::TIME32: + return WriteArrowTimes<::arrow::Time32Type>(def_levels, rep_levels, num_levels, + array, ctx, maybe_parent_nulls); + case ::arrow::Type::TIME64: + return WriteArrowTimes<::arrow::Time64Type>(def_levels, rep_levels, num_levels, + array, ctx, maybe_parent_nulls); WRITE_SERIALIZE_CASE(DECIMAL32) WRITE_SERIALIZE_CASE(DECIMAL64) WRITE_SERIALIZE_CASE(DECIMAL128) @@ -2441,7 +2527,12 @@ Status TypedColumnWriterImpl::WriteArrowDense( WRITE_ZERO_COPY_CASE(INT64) WRITE_SERIALIZE_CASE(UINT32) WRITE_SERIALIZE_CASE(UINT64) - WRITE_ZERO_COPY_CASE(TIME64) + case ::arrow::Type::TIME32: + return WriteArrowTimes<::arrow::Time32Type>(def_levels, rep_levels, num_levels, + array, ctx, maybe_parent_nulls); + case ::arrow::Type::TIME64: + return WriteArrowTimes<::arrow::Time64Type>(def_levels, rep_levels, num_levels, + array, ctx, maybe_parent_nulls); WRITE_ZERO_COPY_CASE(DURATION) WRITE_SERIALIZE_CASE(DECIMAL32) WRITE_SERIALIZE_CASE(DECIMAL64) diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index 89f2be42bffd..91e64f8d8e79 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -1392,24 +1392,26 @@ class PARQUET_EXPORT ArrowWriterProperties { return this; } - /// \brief Coerce all timestamps to the specified time unit. + /// \brief Coerce all timestamps and times to the specified time unit. /// \param unit time unit to truncate to. - /// For Parquet versions 1.0 and 2.4, nanoseconds are casted to microseconds. + /// Parquet versions 1.0 and 2.4 support milliseconds and microseconds. + /// Parquet version 2.6 also supports nanoseconds. Builder* coerce_timestamps(::arrow::TimeUnit::type unit) { coerce_timestamps_enabled_ = true; coerce_timestamps_unit_ = unit; return this; } - /// \brief Allow loss of data when truncating timestamps. + /// \brief Allow loss of data when explicitly coercing timestamps or times. /// /// This is disallowed by default and an error will be returned. + /// Only applies when coerce_timestamps is explicitly enabled. Builder* allow_truncated_timestamps() { truncated_timestamps_allowed_ = true; return this; } - /// \brief Disallow loss of data when truncating timestamps (default). + /// \brief Disallow loss of data when truncating timestamps or times (default). Builder* disallow_truncated_timestamps() { truncated_timestamps_allowed_ = false; return this; diff --git a/docs/source/python/parquet/parquet_type_handling.rst b/docs/source/python/parquet/parquet_type_handling.rst index ccf2fa8a109c..81c8e1df63cc 100644 --- a/docs/source/python/parquet/parquet_type_handling.rst +++ b/docs/source/python/parquet/parquet_type_handling.rst @@ -107,14 +107,20 @@ until PyArrow version ``21.0.0``), set ``arrow_extensions_enabled=False``. (`geoarrow-pyarrow `_) and import ``geoarrow.pyarrow`` module. -Storing timestamps ------------------- +Storing timestamps and times +---------------------------- + +The ``coerce_timestamps`` and ``allow_truncated_timestamps`` options also apply +to Arrow ``time32`` and ``time64`` values. For versions 1.0 and 2.4, nanosecond +times are converted to microseconds by default. Lossy conversion raises an +exception unless both an explicit coercion target and permission to truncate +are supplied. Millisecond times are stored as INT32; microsecond and nanosecond +times are stored as INT64. Some Parquet readers may only support timestamps stored in millisecond (``'ms'``) or microsecond (``'us'``) resolution. Since pandas uses nanoseconds -to represent timestamps, this can occasionally be a nuisance. By default -(when writing version 1.0 Parquet files), the nanoseconds will be cast to -microseconds (``'us'``). +to represent timestamps, this can occasionally be a nuisance. When writing +version 1.0 or 2.4 Parquet files, nanoseconds are cast to microseconds (``'us'``). In addition, we provide the ``coerce_timestamps`` option to allow you to select the desired resolution: @@ -132,8 +138,8 @@ an exception will be raised. This can be suppressed by passing >>> pq.write_table(table, 'example.parquet', coerce_timestamps='ms', ... allow_truncated_timestamps=True) -Timestamps with nanoseconds can be stored without casting when using the -more recent Parquet format version 2.6: +Timestamps and time-of-day values with nanoseconds can be stored without casting +when using the more recent Parquet format version 2.6: .. code-block:: python diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index 4de88fe22997..f24b28b6b10d 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -765,7 +765,7 @@ def _sanitize_table(table, new_schema, flavor): Parquet implementations, so version='1.0' is likely the choice that maximizes file compatibility. UINT32 and some logical types are only available with version '2.4'. - Nanosecond timestamps are only available with version '2.6'. + Nanosecond timestamps and times are only available with version '2.6'. Other features such as compression algorithms or the new serialized data page format must be enabled separately (see 'compression' and 'data_page_version'). @@ -787,7 +787,8 @@ def _sanitize_table(table, new_schema, flavor): Write timestamps to INT96 Parquet format. Defaults to False unless enabled by flavor argument. This take priority over the coerce_timestamps option. coerce_timestamps : str, default None - Cast timestamps to a particular resolution. If omitted, defaults are chosen + Cast timestamps and time-of-day values to a particular resolution. + If omitted, defaults are chosen depending on `version`. For ``version='1.0'`` and ``version='2.4'``, nanoseconds are cast to microseconds ('us'), while for ``version='2.6'`` (the default), they are written natively without loss @@ -797,9 +798,10 @@ def _sanitize_table(table, new_schema, flavor): unless ``allow_truncated_timestamps=True`` is given. Valid values: {None, 'ms', 'us'} allow_truncated_timestamps : bool, default False - Allow loss of data when coercing timestamps to a particular - resolution. E.g. if microsecond or nanosecond data is lost when coercing to - 'ms', do not raise an exception. Passing ``allow_truncated_timestamp=True`` + Allow loss of data when coercing timestamps or time-of-day values + to a particular resolution. E.g. if microsecond or nanosecond data is lost + when coercing to 'ms', do not raise an exception. Passing + ``allow_truncated_timestamp=True`` will NOT result in the truncation exception being ignored unless ``coerce_timestamps`` is not None. data_page_size : int, default None diff --git a/python/pyarrow/tests/parquet/test_datetime.py b/python/pyarrow/tests/parquet/test_datetime.py index b89fd97cb91e..2a09a416e4e4 100644 --- a/python/pyarrow/tests/parquet/test_datetime.py +++ b/python/pyarrow/tests/parquet/test_datetime.py @@ -158,6 +158,116 @@ def test_coerce_timestamps_truncated(tempdir): tm.assert_frame_equal(df_expected, df_ms) +@pytest.mark.parametrize("version, source_unit, target_unit, expected_unit", [ + ("1.0", "ns", None, "us"), + ("2.4", "ns", None, "us"), + ("2.6", "ns", None, "ns"), + ("2.6", "s", None, "ms"), + ("2.6", "ns", "us", "us"), + ("1.0", "ns", "ms", "ms"), + ("2.6", "ms", "us", "us"), + ("2.6", "us", "ms", "ms"), + ("2.6", "s", "us", "us"), +]) +@pytest.mark.parametrize("store_schema", [False, True]) +def test_coerce_times(version, source_unit, target_unit, expected_unit, + store_schema): + units_per_second = {"s": 1, "ms": 1000, "us": 1000000, "ns": 1000000000} + source_type = (pa.time32(source_unit) if source_unit in ("s", "ms") + else pa.time64(source_unit)) + array = pa.array([0, 40984 * units_per_second[source_unit], None], + source_type) + expected_type = (pa.time32(expected_unit) if expected_unit == "ms" + else pa.time64(expected_unit)) + expected = pa.array([0, 40984 * units_per_second[expected_unit], None], + expected_type) + buffer = io.BytesIO() + pq.write_table(pa.table({"time": array}), buffer, version=version, + coerce_timestamps=target_unit, store_schema=store_schema) + buffer.seek(0) + file = pq.ParquetFile(buffer) + physical_type = "INT32" if expected_unit == "ms" else "INT64" + assert file.schema.column(0).physical_type == physical_type + assert file.schema.to_arrow_schema().field(0).type == expected_type + assert file.read().column(0).combine_chunks().equals(expected) + + +@pytest.mark.parametrize("target_unit, value", [("us", 1001), ("ms", 1000001)]) +def test_coerce_times_truncated(target_unit, value): + table = pa.table({"time": pa.array([value, None], pa.time64("ns"))}) + with pytest.raises(pa.ArrowInvalid, match="would lose data"): + pq.write_table(table, io.BytesIO(), coerce_timestamps=target_unit) + buffer = io.BytesIO() + pq.write_table(table, buffer, coerce_timestamps=target_unit, + allow_truncated_timestamps=True) + buffer.seek(0) + dtype = pa.time32("ms") if target_unit == "ms" else pa.time64("us") + assert pq.read_table(buffer).column(0).combine_chunks().equals( + pa.array([1, None], dtype)) + + +@pytest.mark.parametrize("version", ["1.0", "2.4"]) +def test_time_default_coercion_rejects_truncation(version): + table = pa.table({"time": pa.array([1001], pa.time64("ns"))}) + with pytest.raises(pa.ArrowInvalid, match="would lose data"): + pq.write_table(table, io.BytesIO(), version=version, + allow_truncated_timestamps=True) + + +@pytest.mark.parametrize("source_type, target_unit, value", [ + (pa.time32("s"), "ms", 2147483647), + (pa.time64("us"), "ms", 2147483648000), +]) +def test_coerce_times_overflow(source_type, target_unit, value): + table = pa.table({"time": pa.array([value], source_type)}) + with pytest.raises(pa.ArrowInvalid, match="overflow"): + pq.write_table(table, io.BytesIO(), coerce_timestamps=target_unit, + allow_truncated_timestamps=True) + + +def test_coerce_times_null_parent(): + child = pa.array([1001, 2000], pa.time64("ns")) + array = pa.StructArray.from_arrays([child], names=["time"], + mask=pa.array([True, False])) + buffer = io.BytesIO() + pq.write_table(pa.table({"s": array}), buffer, coerce_timestamps="us") + buffer.seek(0) + expected_child = pa.array([None, 2], pa.time64("us")) + expected = pa.StructArray.from_arrays([expected_child], names=["time"], + mask=pa.array([True, False])) + assert pq.read_table(buffer).column(0).combine_chunks().equals(expected) + + +@pytest.mark.parametrize("kind", [ + "chunked", "list", "dictionary", "all_null", +]) +def test_coerce_times_array_layouts(kind): + array = pa.array([7, 1000000, None, 2000000], pa.time64("ns")).slice(1) + expected = pa.array([1000, None, 2000], pa.time64("us")) + if kind == "chunked": + array = pa.chunked_array([array.slice(0, 1), array.slice(1)]) + elif kind == "list": + array = pa.array([[1000000, None], None, [2000000]], + pa.list_(pa.time64("ns"))) + expected = pa.array([[1000, None], None, [2000]], + pa.list_(pa.time64("us"))) + elif kind == "dictionary": + array = pa.DictionaryArray.from_arrays( + pa.array([0, None, 1], pa.int8()), + pa.array([1000000, 2000000], pa.time64("ns"))) + elif kind == "all_null": + array = pa.array([None, None, None], pa.time64("ns")) + expected = pa.array([None, None, None], pa.time64("us")) + buffer = io.BytesIO() + pq.write_table(pa.table({"time": array}), buffer, version="2.4", + row_group_size=1) + buffer.seek(0) + result = pq.read_table(buffer).column(0).combine_chunks() + if pa.types.is_dictionary(result.type): + result = result.dictionary_decode() + assert result.equals(expected) + + @pytest.mark.pandas def test_date_time_types(tempdir): t1 = pa.date32()