This is an automated email from the ASF dual-hosted git repository.
yuxia pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 5b7f839 feat(python): add get_primary_keys() method to Schema class
(#436)
5b7f839 is described below
commit 5b7f8393120f6c33a5264f8d62c28b6a23295321
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Mar 11 22:12:26 2026 +0800
feat(python): add get_primary_keys() method to Schema class (#436)
---------
Co-authored-by: xiaohongbo <[email protected]>
---
bindings/python/fluss/__init__.pyi | 1 +
bindings/python/src/metadata.rs | 8 +++++-
bindings/python/test/test_admin.py | 1 +
bindings/python/test/test_schema.py | 37 +++++++++++++++++++++++++
website/docs/user-guide/python/api-reference.md | 2 ++
5 files changed, 48 insertions(+), 1 deletion(-)
diff --git a/bindings/python/fluss/__init__.pyi
b/bindings/python/fluss/__init__.pyi
index 417ac9b..e06e932 100644
--- a/bindings/python/fluss/__init__.pyi
+++ b/bindings/python/fluss/__init__.pyi
@@ -775,6 +775,7 @@ class Schema:
def get_column_names(self) -> List[str]: ...
def get_column_types(self) -> List[str]: ...
def get_columns(self) -> List[Tuple[str, str]]: ...
+ def get_primary_keys(self) -> List[str]: ...
def __str__(self) -> str: ...
class TableDescriptor:
diff --git a/bindings/python/src/metadata.rs b/bindings/python/src/metadata.rs
index d6b122d..02ef121 100644
--- a/bindings/python/src/metadata.rs
+++ b/bindings/python/src/metadata.rs
@@ -220,7 +220,13 @@ impl Schema {
.collect()
}
- // TODO: support primaryKey
+ /// Get primary key column names, returns empty list if no primary key is
defined
+ fn get_primary_keys(&self) -> Vec<String> {
+ self.__schema
+ .primary_key()
+ .map(|pk| pk.column_names().to_vec())
+ .unwrap_or_default()
+ }
fn __str__(&self) -> String {
format!("Schema: columns={:?}", self.get_columns())
diff --git a/bindings/python/test/test_admin.py
b/bindings/python/test/test_admin.py
index e2f4343..646248d 100644
--- a/bindings/python/test/test_admin.py
+++ b/bindings/python/test/test_admin.py
@@ -82,6 +82,7 @@ async def test_create_table(admin):
),
primary_keys=["id"],
)
+ assert schema.get_primary_keys() == ["id"]
table_descriptor = fluss.TableDescriptor(
schema,
diff --git a/bindings/python/test/test_schema.py
b/bindings/python/test/test_schema.py
new file mode 100644
index 0000000..a72d933
--- /dev/null
+++ b/bindings/python/test/test_schema.py
@@ -0,0 +1,37 @@
+# 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.
+
+"""Unit tests for Schema (no cluster required)."""
+
+import pyarrow as pa
+
+import fluss
+
+
+def test_get_primary_keys():
+ fields = pa.schema([
+ pa.field("id", pa.int32()),
+ pa.field("name", pa.string()),
+ ])
+
+ schema_with_pk = fluss.Schema(fields, primary_keys=["id"])
+ assert schema_with_pk.get_primary_keys() == ["id"]
+
+ schema_without_pk = fluss.Schema(fields)
+ assert schema_without_pk.get_primary_keys() == []
+
+
diff --git a/website/docs/user-guide/python/api-reference.md
b/website/docs/user-guide/python/api-reference.md
index e9113b6..fef10a8 100644
--- a/website/docs/user-guide/python/api-reference.md
+++ b/website/docs/user-guide/python/api-reference.md
@@ -241,6 +241,8 @@ for record in scan_records:
| `Schema(schema: pa.Schema, primary_keys=None)` | Create from PyArrow schema |
| `.get_column_names() -> list[str]` | Get column names |
| `.get_column_types() -> list[str]` | Get column type names |
+| `.get_columns() -> list[tuple[str, str]]` | Get `(name, type)` pairs |
+| `.get_primary_keys() -> list[str]` | Get primary key columns |
## `TableDescriptor`
