From ffca82620688524731578aa9b8f3360bb3b3698f Mon Sep 17 00:00:00 2001 From: charlesdong1991 Date: Sat, 2 May 2026 15:32:15 +0200 Subject: [PATCH] Add nullability support in python --- Cargo.lock | 118 ++++++++++++++++++ bindings/python/src/metadata.rs | 2 +- bindings/python/src/utils.rs | 34 +++-- bindings/python/test/test_schema.py | 47 +++++++ .../docs/user-guide/python/api-reference.md | 6 +- website/docs/user-guide/python/data-types.md | 25 ++++ 6 files changed, 217 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ac5b27f8..72a86180 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -746,6 +746,15 @@ dependencies = [ "rustc_version", ] +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -975,6 +984,12 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "endian-type" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" + [[package]] name = "equivalent" version = "1.0.2" @@ -1104,6 +1119,8 @@ dependencies = [ "jiff", "linked-hash-map", "log", + "metrics", + "metrics-util", "opendal", "ordered-float", "parking_lot", @@ -1403,6 +1420,9 @@ name = "hashbrown" version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "foldhash 0.2.0", +] [[package]] name = "heck" @@ -2054,6 +2074,37 @@ dependencies = [ "autocfg", ] +[[package]] +name = "metrics" +version = "0.24.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff56c2e7dce6bd462e3b8919986a617027481b1dcc703175b58cf9dd98a2f071" +dependencies = [ + "portable-atomic", + "rapidhash", +] + +[[package]] +name = "metrics-util" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e56997f084e57b045edf17c3ed8ba7f9f779c670df8206dfd1c736f4c02dc4a" +dependencies = [ + "aho-corasick", + "crossbeam-epoch", + "crossbeam-utils", + "hashbrown 0.16.1", + "indexmap 2.13.1", + "metrics", + "ordered-float", + "quanta", + "radix_trie", + "rand 0.9.2", + "rand_xoshiro", + "rapidhash", + "sketches-ddsketch", +] + [[package]] name = "mime" version = "0.3.17" @@ -2077,6 +2128,15 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" +[[package]] +name = "nibble_vec" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43" +dependencies = [ + "smallvec", +] + [[package]] name = "num" version = "0.4.3" @@ -2546,6 +2606,21 @@ dependencies = [ "cc", ] +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + [[package]] name = "quick-xml" version = "0.37.5" @@ -2648,6 +2723,16 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "radix_trie" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" +dependencies = [ + "endian-type", + "nibble_vec", +] + [[package]] name = "rand" version = "0.8.5" @@ -2709,6 +2794,33 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_xoshiro" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f703f4665700daf5512dcca5f43afa6af89f09db47fb56be587f80636bda2d41" +dependencies = [ + "rand_core 0.9.5", +] + +[[package]] +name = "rapidhash" +version = "4.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e48930979c155e2f33aa36ab3119b5ee81332beb6482199a8ecd6029b80b59" +dependencies = [ + "rustversion", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3214,6 +3326,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "sketches-ddsketch" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6f73aeb92d671e0cc4dca167e59b2deb6387c375391bc99ee743f326994a2b" + [[package]] name = "slab" version = "0.4.12" diff --git a/bindings/python/src/metadata.rs b/bindings/python/src/metadata.rs index 02ef121d..7b6129a4 100644 --- a/bindings/python/src/metadata.rs +++ b/bindings/python/src/metadata.rs @@ -165,7 +165,7 @@ impl Schema { let mut builder = fcore::metadata::Schema::builder(); for field in arrow_schema.fields() { - let fluss_data_type = crate::utils::Utils::arrow_type_to_fluss_type(field.data_type())?; + let fluss_data_type = crate::utils::Utils::arrow_field_to_fluss_type(field)?; builder = builder.column(field.name(), fluss_data_type); if let Some(comment) = field.metadata().get("comment") { diff --git a/bindings/python/src/utils.rs b/bindings/python/src/utils.rs index 76b95f7a..5efcf5e7 100644 --- a/bindings/python/src/utils.rs +++ b/bindings/python/src/utils.rs @@ -36,14 +36,14 @@ impl Utils { }) } - /// Convert Arrow DataType to Fluss DataType - pub fn arrow_type_to_fluss_type( - arrow_type: &arrow::datatypes::DataType, + /// Convert an Arrow Field to a Fluss DataType, preserving nullability. + pub fn arrow_field_to_fluss_type( + field: &arrow::datatypes::Field, ) -> PyResult { use arrow::datatypes::DataType as ArrowDataType; use fcore::metadata::DataTypes; - let fluss_type = match arrow_type { + let fluss_type = match field.data_type() { ArrowDataType::Boolean => DataTypes::boolean(), ArrowDataType::Int8 => DataTypes::tinyint(), ArrowDataType::Int16 => DataTypes::smallint(), @@ -95,23 +95,29 @@ impl Utils { ArrowDataType::Decimal128(precision, scale) => { DataTypes::decimal(*precision as u32, *scale as u32) } - ArrowDataType::List(field) => { - let element_type = Utils::arrow_type_to_fluss_type(field.data_type())?; + ArrowDataType::List(element_field) => { + let element_type = Utils::arrow_field_to_fluss_type(element_field)?; DataTypes::array(element_type) } - _ => { + other => { return Err(FlussError::new_err(format!( - "Unsupported Arrow data type: {arrow_type:?}" + "Unsupported Arrow data type: {other:?}" ))); } }; - Ok(fluss_type) + if field.is_nullable() { + Ok(fluss_type) + } else { + Ok(fluss_type.as_non_nullable()) + } } - /// Convert Fluss DataType to string representation + /// Convert Fluss DataType to string representation, appending " NOT NULL" + /// for non-nullable types (matches Java's `withNullability` and Rust core's + /// `Display` impl). pub fn datatype_to_string(data_type: &fcore::metadata::DataType) -> String { - match data_type { + let type_str = match data_type { fcore::metadata::DataType::Boolean(_) => "boolean".to_string(), fcore::metadata::DataType::TinyInt(_) => "tinyint".to_string(), fcore::metadata::DataType::SmallInt(_) => "smallint".to_string(), @@ -171,6 +177,12 @@ impl Utils { .collect(); format!("row<{}>", fields.join(", ")) } + }; + + if data_type.is_nullable() { + type_str + } else { + format!("{type_str} NOT NULL") } } diff --git a/bindings/python/test/test_schema.py b/bindings/python/test/test_schema.py index ab2d1ab4..dfd9cf56 100644 --- a/bindings/python/test/test_schema.py +++ b/bindings/python/test/test_schema.py @@ -48,3 +48,50 @@ def test_schema_with_array(): assert schema.get_column_types() == ["int", "array"] +def test_nullable_fields(): + fields = pa.schema( + [ + pa.field("id", pa.int32(), nullable=False), + pa.field("name", pa.string()), + ] + ) + schema = fluss.Schema(fields) + assert schema.get_column_types() == ["int NOT NULL", "string"] + assert schema.get_columns() == [("id", "int NOT NULL"), ("name", "string")] + + +def test_pk_forces_non_nullable(): + fields = pa.schema( + [ + pa.field("id", pa.int32()), + pa.field("name", pa.string()), + ] + ) + schema = fluss.Schema(fields, primary_keys=["id"]) + types = schema.get_column_types() + assert types[0] == "int NOT NULL" + assert types[1] == "string" + + +def test_nested_list_nullability(): + fields = pa.schema( + [ + pa.field( + "tags", + pa.list_(pa.field("item", pa.string(), nullable=False)), + ), + pa.field("ids", pa.list_(pa.int32()), nullable=False), + pa.field( + "strict_ids", + pa.list_(pa.field("item", pa.int32(), nullable=False)), + nullable=False, + ), + ] + ) + schema = fluss.Schema(fields) + types = schema.get_column_types() + assert types[0] == "array" + assert types[1] == "array NOT NULL" + assert types[2] == "array NOT NULL" + + diff --git a/website/docs/user-guide/python/api-reference.md b/website/docs/user-guide/python/api-reference.md index 73b9a8fe..92f2a4d1 100644 --- a/website/docs/user-guide/python/api-reference.md +++ b/website/docs/user-guide/python/api-reference.md @@ -242,10 +242,10 @@ for record in scan_records: | Method | Description | |------------------------------------------------|----------------------------| -| `Schema(schema: pa.Schema, primary_keys=None)` | Create from PyArrow schema | +| `Schema(schema: pa.Schema, primary_keys=None)` | Create from PyArrow schema. Field nullability (`pa.field(..., nullable=False)`) is preserved. | | `.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_column_types() -> list[str]` | Get column type names. Non-nullable types include a `" NOT NULL"` suffix (e.g., `"int NOT NULL"`). | +| `.get_columns() -> list[tuple[str, str]]` | Get `(name, type)` pairs. Type strings follow the same nullability formatting as `.get_column_types()`. | | `.get_primary_keys() -> list[str]` | Get primary key columns | ## `TableDescriptor` diff --git a/website/docs/user-guide/python/data-types.md b/website/docs/user-guide/python/data-types.md index c0acb4c7..0908b84c 100644 --- a/website/docs/user-guide/python/data-types.md +++ b/website/docs/user-guide/python/data-types.md @@ -21,6 +21,31 @@ The Python client uses PyArrow types for schema definitions: All Python native types (`date`, `time`, `datetime`, `Decimal`) work when appending rows via dicts. +## Nullability + +PyArrow field nullability is preserved when constructing Fluss schemas. By default, fields are nullable. Use `nullable=False` on `pa.field()` to create a `NOT NULL` column: + +```python +schema = pa.schema([ + pa.field("id", pa.int32(), nullable=False), + pa.field("name", pa.string()), # nullable by default +]) +fluss_schema = fluss.Schema(schema) +fluss_schema.get_column_types() # ["int NOT NULL", "string"] +``` + +Primary key columns are automatically forced `NOT NULL` regardless of the PyArrow field setting. + +For nested types, element nullability is also preserved: + +```python +schema = pa.schema([ + pa.field("tags", pa.list_(pa.field("item", pa.string(), nullable=False))), +]) +fluss_schema = fluss.Schema(schema) +fluss_schema.get_column_types() # ["array"] +``` + ## Writing Data Rows can be dicts, lists, or tuples: