From bb0f973bed5a4c81bd4305c56920baaebd4298af Mon Sep 17 00:00:00 2001 From: zhangzhibiao Date: Wed, 24 Jun 2026 13:08:47 +0800 Subject: [PATCH] support flink ts --- velox/connectors/pulsar/PulsarDataSource.cpp | 2 +- velox/connectors/pulsar/PulsarDataSource.h | 2 +- .../tests/PulsarConnectorIntegrationTest.cpp | 6 +- velox/type/Type.cpp | 64 +++++++++++++++++++ velox/type/Type.h | 37 +++++++++++ velox/type/tests/TypeTest.cpp | 26 ++++++++ 6 files changed, 132 insertions(+), 5 deletions(-) diff --git a/velox/connectors/pulsar/PulsarDataSource.cpp b/velox/connectors/pulsar/PulsarDataSource.cpp index 3ea84f3588d..14237616d5c 100644 --- a/velox/connectors/pulsar/PulsarDataSource.cpp +++ b/velox/connectors/pulsar/PulsarDataSource.cpp @@ -300,7 +300,7 @@ std::optional PulsarDataSource::next( return res; } -std::vector PulsarDataSource::checkpointState() { +std::vector PulsarDataSource::snapshotState(int64_t) { if (checkpointStartMessageId_.empty()) { return {}; } diff --git a/velox/connectors/pulsar/PulsarDataSource.h b/velox/connectors/pulsar/PulsarDataSource.h index 2f693901874..6963b711e18 100644 --- a/velox/connectors/pulsar/PulsarDataSource.h +++ b/velox/connectors/pulsar/PulsarDataSource.h @@ -60,7 +60,7 @@ class PulsarDataSource : public DataSource { return completedRows_; } - std::vector checkpointState() override; + std::vector snapshotState(int64_t checkpointId) override; std::vector commit(int64_t id) override; diff --git a/velox/connectors/pulsar/tests/PulsarConnectorIntegrationTest.cpp b/velox/connectors/pulsar/tests/PulsarConnectorIntegrationTest.cpp index 7096e9851c0..550c7047562 100644 --- a/velox/connectors/pulsar/tests/PulsarConnectorIntegrationTest.cpp +++ b/velox/connectors/pulsar/tests/PulsarConnectorIntegrationTest.cpp @@ -242,7 +242,7 @@ TEST(PulsarConnectorIntegrationTest, checkpointStateRecordsLastMessageId) { auto connectorConfig = makeRawConfig(serviceUrl, topic, subscription); auto source = createRawDataSource( pool, connectorConfig, connectorId, serviceUrl, topic, subscription); - ASSERT_TRUE(source->checkpointState().empty()); + ASSERT_TRUE(source->snapshotState(1).empty()); std::vector<::pulsar::MessageId> messageIds; produceRawMessages(serviceUrl, topic, messageIds); @@ -253,7 +253,7 @@ TEST(PulsarConnectorIntegrationTest, checkpointStateRecordsLastMessageId) { ASSERT_NE(resultVector.value(), nullptr); ASSERT_EQ(resultVector.value()->size(), 2); - const auto checkpointState = source->checkpointState(); + const auto checkpointState = source->snapshotState(1); ASSERT_EQ(checkpointState.size(), 1); const auto checkpointSplit = PulsarConnectorSplit::create(folly::parseJson(checkpointState[0])); @@ -308,7 +308,7 @@ TEST(PulsarConnectorIntegrationTest, checkpointModeAcksOnCommit) { auto stats = source->runtimeStats(); ASSERT_EQ(stats.at("pulsarAcknowledgedMessages").value, 0); - const auto checkpointState = source->checkpointState(); + const auto checkpointState = source->snapshotState(1); ASSERT_EQ(checkpointState.size(), 1); const auto committed = source->commit(1); ASSERT_EQ(committed, checkpointState); diff --git a/velox/type/Type.cpp b/velox/type/Type.cpp index f7791200f2e..e210ed5488f 100644 --- a/velox/type/Type.cpp +++ b/velox/type/Type.cpp @@ -174,6 +174,12 @@ TypePtr Type::create(const folly::dynamic& obj) { if (isDecimalName(typeName)) { return DECIMAL(obj["precision"].asInt(), obj["scale"].asInt()); } + if (typeName == "FLINK_TIMESTAMP") { + return FLINK_TIMESTAMP(obj["precision"].asInt()); + } + if (typeName == "FLINK_TIMESTAMP_LTZ") { + return FLINK_TIMESTAMP_LTZ(obj["precision"].asInt()); + } // Checks if 'typeName' specifies a custom type. if (customTypeExists(typeName)) { std::vector params; @@ -921,6 +927,40 @@ VELOX_DEFINE_SCALAR_ACCESSOR(VARBINARY); #undef VELOX_DEFINE_SCALAR_ACCESSOR +FlinkTimestampType::FlinkTimestampType(int precision, bool localZoned) + : TimestampType(), + localZoned_(localZoned), + parameters_{TypeParameter(precision)} { + VELOX_CHECK_GE(precision, 0, "Timestamp precision must be at least 0"); + VELOX_CHECK_LE(precision, 9, "Timestamp precision must not exceed 9"); +} + +bool FlinkTimestampType::equivalent(const Type& other) const { + if (!Type::hasSameTypeId(other)) { + return false; + } + const auto& otherTimestamp = static_cast(other); + return precision() == otherTimestamp.precision() && + localZoned_ == otherTimestamp.localZoned_; +} + +folly::dynamic FlinkTimestampType::serialize() const { + auto obj = TimestampType::serialize(); + obj["type"] = name(); + obj["precision"] = precision(); + return obj; +} + +std::shared_ptr FLINK_TIMESTAMP(int precision) { + return std::make_shared( + precision, false /* localZoned */); +} + +std::shared_ptr FLINK_TIMESTAMP_LTZ(int precision) { + return std::make_shared( + precision, true /* localZoned */); +} + TypePtr UNKNOWN() { return TypeFactory::create(); } @@ -1333,6 +1373,28 @@ class FunctionParametricType { } }; +class FlinkTimestampParametricType { + public: + static TypePtr create(const std::vector& parameters) { + VELOX_USER_CHECK_EQ(1, parameters.size()); + VELOX_USER_CHECK(parameters[0].kind == TypeParameterKind::kLongLiteral); + VELOX_USER_CHECK(parameters[0].longLiteral.has_value()); + + return FLINK_TIMESTAMP(parameters[0].longLiteral.value()); + } +}; + +class FlinkTimestampLtzParametricType { + public: + static TypePtr create(const std::vector& parameters) { + VELOX_USER_CHECK_EQ(1, parameters.size()); + VELOX_USER_CHECK(parameters[0].kind == TypeParameterKind::kLongLiteral); + VELOX_USER_CHECK(parameters[0].longLiteral.has_value()); + + return FLINK_TIMESTAMP_LTZ(parameters[0].longLiteral.value()); + } +}; + using ParametricTypeMap = std::unordered_map< std::string, std::function& parameters)>>; @@ -1344,6 +1406,8 @@ const ParametricTypeMap& parametricBuiltinTypes() { {"MAP", MapParametricType::create}, {"ROW", RowParametricType::create}, {"FUNCTION", FunctionParametricType::create}, + {"FLINK_TIMESTAMP", FlinkTimestampParametricType::create}, + {"FLINK_TIMESTAMP_LTZ", FlinkTimestampLtzParametricType::create}, }; return kTypes; } diff --git a/velox/type/Type.h b/velox/type/Type.h index 266b5df6cbf..4e538e01e0b 100644 --- a/velox/type/Type.h +++ b/velox/type/Type.h @@ -1254,6 +1254,43 @@ using TimestampType = ScalarType; using VarcharType = ScalarType; using VarbinaryType = ScalarType; +class FlinkTimestampType : public TimestampType { + public: + explicit FlinkTimestampType(int precision, bool localZoned = false); + + uint8_t precision() const { + return parameters_[0].longLiteral.value(); + } + + bool localZoned() const { + return localZoned_; + } + + const char* name() const override { + return localZoned_ ? "FLINK_TIMESTAMP_LTZ" : "FLINK_TIMESTAMP"; + } + + std::string toString() const override { + return fmt::format( + "{}({})", localZoned_ ? "TIMESTAMP_LTZ" : "TIMESTAMP", precision()); + } + + bool equivalent(const Type& other) const override; + + folly::dynamic serialize() const override; + + const std::vector& parameters() const override { + return parameters_; + } + + private: + const bool localZoned_; + const std::vector parameters_; +}; + +std::shared_ptr FLINK_TIMESTAMP(int precision); +std::shared_ptr FLINK_TIMESTAMP_LTZ(int precision); + constexpr long kMillisInSecond = 1000; constexpr long kMillisInMinute = 60 * kMillisInSecond; constexpr long kMillisInHour = 60 * kMillisInMinute; diff --git a/velox/type/tests/TypeTest.cpp b/velox/type/tests/TypeTest.cpp index f315c149045..dc0bb9c097a 100644 --- a/velox/type/tests/TypeTest.cpp +++ b/velox/type/tests/TypeTest.cpp @@ -58,6 +58,32 @@ TEST(TypeTest, constructorThrow) { "[names: {'a', 'b'}, types: {VARCHAR, NULL}]"); } +TEST(TypeTest, flinkTimestampSerde) { + Type::registerSerDe(); + + auto timestamp = + velox::ISerializable::deserialize(folly::dynamic::object( + "name", "Type")("type", "FLINK_TIMESTAMP")("precision", 3)); + ASSERT_EQ(timestamp->kind(), TypeKind::TIMESTAMP); + ASSERT_STREQ(timestamp->name(), "FLINK_TIMESTAMP"); + ASSERT_EQ(timestamp->toString(), "TIMESTAMP(3)"); + ASSERT_EQ(timestamp->serialize()["precision"].asInt(), 3); + + auto timestampLtz = + velox::ISerializable::deserialize(folly::dynamic::object( + "name", "Type")("type", "FLINK_TIMESTAMP_LTZ")("precision", 9)); + ASSERT_EQ(timestampLtz->kind(), TypeKind::TIMESTAMP); + ASSERT_STREQ(timestampLtz->name(), "FLINK_TIMESTAMP_LTZ"); + ASSERT_EQ(timestampLtz->toString(), "TIMESTAMP_LTZ(9)"); + ASSERT_EQ(timestampLtz->serialize()["precision"].asInt(), 9); + ASSERT_NE(*timestamp, *timestampLtz); + + auto timestampCopy = + velox::ISerializable::deserialize(timestamp->serialize()); + ASSERT_EQ(*timestamp, *timestampCopy); + ASSERT_EQ(timestampCopy->toString(), "TIMESTAMP(3)"); +} + TEST(TypeTest, array) { const auto arrayType = ARRAY(ARRAY(ARRAY(INTEGER()))); ASSERT_EQ("ARRAY>>", arrayType->toString());