feat: sort order

Signed-off-by: Junwang Zhao <[email protected]>

Author: zhjwpku

src/iceberg/CMakeLists.txt

@@ -25,6 +25,8 @@ set(ICEBERG_SOURCES
     schema_internal.cc
     partition_field.cc
     partition_spec.cc
+    sort_field.cc
+    sort_order.cc
     transform.cc
     type.cc)
 

src/iceberg/sort_field.cc

@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "iceberg/sort_field.h"
+
+#include <format>
+
+#include "iceberg/transform.h"
+#include "iceberg/type.h"
+#include "iceberg/util/formatter.h"
+
+namespace iceberg {
+
+namespace {
+/// \brief Get the relative sort direction name
+constexpr std::string_view ToString(SortDirection direction) {
+  switch (direction) {
+    case SortDirection::kAscending:
+      return "asc";
+    case SortDirection::kDescending:
+      return "desc";
+    default:
+      return "invalid";
+  }
+}
+
+/// \brief Get the relative null order name
+constexpr std::string_view ToString(NullOrder null_order) {
+  switch (null_order) {
+    case NullOrder::kFirst:
+      return "nulls-first";
+    case NullOrder::kLast:
+      return "nulls-last";
+    default:
+      return "invalid";
+  }
+}
+}  // namespace
+
+SortField::SortField(int32_t source_id, std::shared_ptr<TransformFunction> transform,
+                     SortDirection sort_direction, NullOrder null_order)
+    : source_id_(source_id),
+      transform_(std::move(transform)),
+      sort_direction_(sort_direction),
+      null_order_(null_order) {}
+
+int32_t SortField::source_id() const { return source_id_; }
+
+std::shared_ptr<TransformFunction> const& SortField::transform() const {
+  return transform_;
+}
+
+SortDirection SortField::sort_direction() const { return sort_direction_; }
+
+NullOrder SortField::null_order() const { return null_order_; }
+
+std::string SortField::ToString() const {
+  return std::format(
+      "SortField(source_id={}, transform={}, sort_direction={}, null_order={})",
+      source_id_, *transform_, iceberg::ToString(sort_direction_),
+      iceberg::ToString(null_order_));
+}
+
+bool SortField::Equals(const SortField& other) const {
+  return source_id_ == other.source_id_ && *transform_ == *other.transform_ &&
+         sort_direction_ == other.sort_direction_ && null_order_ == other.null_order_;
+}
+
+}  // namespace iceberg

src/iceberg/sort_field.h

@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#pragma once
+
+/// \file iceberg/sort_field.h
+/// A sort field in a sort order
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include "iceberg/iceberg_export.h"
+#include "iceberg/type_fwd.h"
+#include "iceberg/util/formattable.h"
+
+namespace iceberg {
+
+/// \brief Sort direction in a partition, either ascending or descending
+enum class SortDirection {
+  /// Ascending
+  kAscending,
+  /// Descending
+  kDescending,
+};
+
+enum class NullOrder {
+  /// Nulls are sorted first
+  kFirst,
+  /// Nulls are sorted last
+  kLast,
+};
+
+/// \brief a field with its transform.
+class ICEBERG_EXPORT SortField : public util::Formattable {
+ public:
+  /// \brief Construct a field.
+  /// \param[in] source_id The source field ID.
+  /// \param[in] transform The transform function.
+  /// \param[in] sort_direction The sort direction.
+  /// \param[in] null_order The null order.
+  SortField(int32_t source_id, std::shared_ptr<TransformFunction> transform,
+            SortDirection sort_direction, NullOrder null_order);
+
+  /// \brief Get the source field ID.
+  int32_t source_id() const;
+
+  /// \brief Get the transform type.
+  const std::shared_ptr<TransformFunction>& transform() const;
+
+  /// \brief Get the sort direction.
+  SortDirection sort_direction() const;
+
+  /// \brief Get the null order.
+  NullOrder null_order() const;
+
+  std::string ToString() const override;
+
+  friend bool operator==(const SortField& lhs, const SortField& rhs) {
+    return lhs.Equals(rhs);
+  }
+
+  friend bool operator!=(const SortField& lhs, const SortField& rhs) {
+    return !(lhs == rhs);
+  }
+
+ private:
+  /// \brief Compare two fields for equality.
+  [[nodiscard]] bool Equals(const SortField& other) const;
+
+  int32_t source_id_;
+  std::shared_ptr<TransformFunction> transform_;
+  SortDirection sort_direction_;
+  NullOrder null_order_;
+};
+
+}  // namespace iceberg

src/iceberg/sort_order.cc

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "iceberg/sort_order.h"
+
+#include <format>
+
+#include "iceberg/util/formatter.h"
+
+namespace iceberg {
+
+SortOrder::SortOrder(int64_t order_id, std::vector<SortField> fields)
+    : order_id_(order_id), fields_(std::move(fields)) {}
+
+int64_t SortOrder::order_id() const { return order_id_; }
+
+std::span<const SortField> SortOrder::fields() const { return fields_; }
+
+std::string SortOrder::ToString() const {
+  std::string repr = std::format("sort_order[order_id<{}>,\n", order_id_);
+  for (const auto& field : fields_) {
+    std::format_to(std::back_inserter(repr), "  {}\n", field);
+  }
+  repr += "]";
+  return repr;
+}
+
+bool SortOrder::Equals(const SortOrder& other) const {
+  return order_id_ == other.order_id_ && fields_ == other.fields_;
+}
+
+}  // namespace iceberg

src/iceberg/sort_order.h

@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <span>
+#include <vector>
+
+#include "iceberg/iceberg_export.h"
+#include "iceberg/sort_field.h"
+#include "iceberg/util/formattable.h"
+
+namespace iceberg {
+
+/// \brief A sort order for a table
+///
+/// A sort order is defined by a sort order id and a list of sort fields.
+/// The order of the sort fields within the list defines the order in which the sort is
+/// applied to the data.
+class ICEBERG_EXPORT SortOrder : public util::Formattable {
+ public:
+  SortOrder(int64_t order_id, std::vector<SortField> fields);
+
+  /// \brief Get the sort order id.
+  int64_t order_id() const;
+
+  /// \brief Get the list of sort fields.
+  std::span<const SortField> fields() const;
+
+  std::string ToString() const override;
+
+  friend bool operator==(const SortOrder& lhs, const SortOrder& rhs) {
+    return lhs.Equals(rhs);
+  }
+
+  friend bool operator!=(const SortOrder& lhs, const SortOrder& rhs) {
+    return !(lhs == rhs);
+  }
+
+ private:
+  /// \brief Compare two sort orders for equality.
+  bool Equals(const SortOrder& other) const;
+
+  int64_t order_id_;
+  std::vector<SortField> fields_;
+};
+
+}  // namespace iceberg

test/CMakeLists.txt

@@ -30,7 +30,9 @@ target_sources(schema_test
                        type_test.cc
                        transform_test.cc
                        partition_field_test.cc
-                       partition_spec_test.cc)
+                       partition_spec_test.cc
+                       sort_field_test.cc
+                       sort_order_test.cc)
 target_link_libraries(schema_test PRIVATE iceberg_static GTest::gtest_main GTest::gmock)
 add_test(NAME schema_test COMMAND schema_test)
 

test/sort_field_test.cc

@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "iceberg/sort_field.h"
+
+#include <format>
+
+#include <gtest/gtest.h>
+
+#include "iceberg/transform.h"
+#include "iceberg/util/formatter.h"
+
+namespace iceberg {
+
+namespace {
+class TestTransformFunction : public TransformFunction {
+ public:
+  TestTransformFunction() : TransformFunction(TransformType::kUnknown) {}
+  expected<ArrowArray, Error> Transform(const ArrowArray& input) override {
+    return unexpected(
+        Error{.kind = ErrorKind::kNotSupported, .message = "test transform function"});
+  }
+};
+
+}  // namespace
+
+TEST(SortFieldTest, Basics) {
+  {
+    const auto transform = std::make_shared<IdentityTransformFunction>();
+    SortField field(1, transform, SortDirection::kAscending, NullOrder::kFirst);
+    EXPECT_EQ(1, field.source_id());
+    EXPECT_EQ(*transform, *field.transform());
+    EXPECT_EQ(SortDirection::kAscending, field.sort_direction());
+    EXPECT_EQ(NullOrder::kFirst, field.null_order());
+    EXPECT_EQ(
+        "SortField(source_id=1, transform=identity, sort_direction=asc, "
+        "null_order=nulls-first)",
+        field.ToString());
+    EXPECT_EQ(
+        "SortField(source_id=1, transform=identity, sort_direction=asc, "
+        "null_order=nulls-first)",
+        std::format("{}", field));
+  }
+}
+
+TEST(SortFieldTest, Equality) {
+  auto test_transform = std::make_shared<TestTransformFunction>();
+  auto identity_transform = std::make_shared<IdentityTransformFunction>();
+
+  SortField field1(1, test_transform, SortDirection::kAscending, NullOrder::kFirst);
+  SortField field2(2, test_transform, SortDirection::kAscending, NullOrder::kFirst);
+  SortField field3(1, identity_transform, SortDirection::kAscending, NullOrder::kFirst);
+  SortField field4(1, test_transform, SortDirection::kDescending, NullOrder::kFirst);
+  SortField field5(1, test_transform, SortDirection::kAscending, NullOrder::kLast);
+
+  ASSERT_EQ(field1, field1);
+  ASSERT_NE(field1, field2);
+  ASSERT_NE(field2, field1);
+  ASSERT_NE(field1, field3);
+  ASSERT_NE(field3, field1);
+  ASSERT_NE(field1, field4);
+  ASSERT_NE(field4, field1);
+  ASSERT_NE(field1, field5);
+  ASSERT_NE(field5, field1);
+}
+}  // namespace iceberg