chore: cleanup and fix partial update in python (#292)

fresh-borzoni · web-flow · commit e1e7ba313ed7 · 2026-02-09T15:14:26.000+08:00
diff --git a/bindings/python/example/example.py b/bindings/python/example/example.py
@@ -569,6 +569,49 @@ async def main():
         print(f"Error during delete: {e}")
         traceback.print_exc()
 
+    # --- Test Partial Update by column names ---
+    print("\n--- Testing Partial Update (by column names) ---")
+    try:
+        partial_writer = pk_table.new_upsert(columns=["user_id", "balance"])
+        handle = partial_writer.upsert({"user_id": 1, "balance": Decimal("9999.99")})
+        await handle.wait()
+        print("Partial update: set balance=9999.99 for user_id=1")
+
+        lookuper = pk_table.new_lookup()
+        result = await lookuper.lookup({"user_id": 1})
+        if result:
+            print(f"Partial update verified:"
+                  f"\n  name={result['name']} (unchanged)"
+                  f"\n  balance={result['balance']} (updated)")
+        else:
+            print("ERROR: Expected to find user_id=1")
+
+    except Exception as e:
+        print(f"Error during partial update by names: {e}")
+        traceback.print_exc()
+
+    # --- Test Partial Update by column indices ---
+    print("\n--- Testing Partial Update (by column indices) ---")
+    try:
+        # Columns: 0=user_id (PK), 1=name — update name only
+        partial_writer_idx = pk_table.new_upsert(column_indices=[0, 1])
+        handle = partial_writer_idx.upsert([1, "Alice Renamed"])
+        await handle.wait()
+        print("Partial update by indices: set name='Alice Renamed' for user_id=1")
+
+        lookuper = pk_table.new_lookup()
+        result = await lookuper.lookup({"user_id": 1})
+        if result:
+            print(f"Partial update by indices verified:"
+                  f"\n  name={result['name']} (updated)"
+                  f"\n  balance={result['balance']} (unchanged)")
+        else:
+            print("ERROR: Expected to find user_id=1")
+
+    except Exception as e:
+        print(f"Error during partial update by indices: {e}")
+        traceback.print_exc()
+
     # Demo: Column projection using builder pattern
     print("\n--- Testing Column Projection ---")
     try:
diff --git a/bindings/python/fluss/__init__.pyi b/bindings/python/fluss/__init__.pyi
@@ -378,15 +378,14 @@ class AppendWriter:
             WriteResultHandle: Ignore for fire-and-forget, or await handle.wait() for acknowledgement.
 
         Supported Types:
-            Currently supports primitive types only:
             - Boolean, TinyInt, SmallInt, Int, BigInt (integers)
             - Float, Double (floating point)
             - String, Char (text)
             - Bytes, Binary (binary data)
+            - Date, Time, Timestamp, TimestampLTZ (temporal)
+            - Decimal (arbitrary precision)
             - Null values
 
-            Temporal types (Date, Timestamp, Decimal) are not yet supported.
-
         Example:
             writer.append({'id': 1, 'name': 'Alice', 'score': 95.5})
             writer.append([1, 'Alice', 95.5])
@@ -712,5 +711,7 @@ class OffsetType:
 
 # Constant for earliest offset (-2)
 EARLIEST_OFFSET: int
+# Constant for latest offset (-1)
+LATEST_OFFSET: int
 
 __version__: str
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
@@ -96,6 +96,7 @@ fn _fluss(m: &Bound<'_, PyModule>) -> PyResult<()> {
 
     // Register constants
     m.add("EARLIEST_OFFSET", fcore::client::EARLIEST_OFFSET)?;
+    m.add("LATEST_OFFSET", fcore::client::LATEST_OFFSET)?;
 
     // Register exception types
     m.add_class::<FlussError>()?;
diff --git a/bindings/python/src/lookup.rs b/bindings/python/src/lookup.rs
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::table::{internal_row_to_dict, python_pk_to_generic_row};
+use crate::table::{internal_row_to_dict, python_to_sparse_generic_row};
 use crate::*;
 use pyo3_async_runtimes::tokio::future_into_py;
 use std::sync::Arc;
@@ -52,7 +52,8 @@ impl Lookuper {
         py: Python<'py>,
         pk: &Bound<'_, PyAny>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        let generic_row = python_pk_to_generic_row(pk, &self.table_info)?;
+        let pk_indices = self.table_info.get_schema().primary_key_indexes();
+        let generic_row = python_to_sparse_generic_row(pk, &self.table_info, &pk_indices)?;
         let inner = self.inner.clone();
         let table_info = self.table_info.clone();
 
diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs
@@ -639,203 +639,100 @@ enum RowInput<'py> {
     List(Bound<'py, pyo3::types::PyList>),
 }
 
-/// Helper function to process sequence types (list/tuple) into datums
-fn process_sequence_to_datums<'a, I>(
-    values: I,
-    len: usize,
-    fields: &[fcore::metadata::DataField],
-) -> PyResult<Vec<fcore::row::Datum<'static>>>
-where
-    I: Iterator<Item = Bound<'a, PyAny>>,
-{
-    if len != fields.len() {
-        return Err(FlussError::new_err(format!(
-            "Expected {} values, got {}",
-            fields.len(),
-            len
-        )));
-    }
-
-    let mut datums = Vec::with_capacity(fields.len());
-    for (i, (field, value)) in fields.iter().zip(values).enumerate() {
-        datums.push(
-            python_value_to_datum(&value, field.data_type()).map_err(|e| {
-                FlussError::new_err(format!("Field '{}' (index {}): {}", field.name(), i, e))
-            })?,
-        );
-    }
-    Ok(datums)
-}
-
-/// Convert Python row (dict/list/tuple) to GenericRow based on schema
+/// Convert Python row (dict/list/tuple) to GenericRow requiring all schema columns.
 pub fn python_to_generic_row(
     row: &Bound<PyAny>,
     table_info: &fcore::metadata::TableInfo,
 ) -> PyResult<fcore::row::GenericRow<'static>> {
-    // Extract with user-friendly error message
-    let row_input: RowInput = row.extract().map_err(|_| {
-        let type_name = row
-            .get_type()
-            .name()
-            .map(|n| n.to_string())
-            .unwrap_or_else(|_| "unknown".to_string());
-        FlussError::new_err(format!(
-            "Row must be a dict, list, or tuple; got {type_name}"
-        ))
-    })?;
-    let schema = table_info.row_type();
-    let fields = schema.fields();
-
-    let datums = match row_input {
-        RowInput::Dict(dict) => {
-            // Strict: reject unknown keys (and also reject non-str keys nicely)
-            for (k, _) in dict.iter() {
-                let key_str = k.extract::<&str>().map_err(|_| {
-                    let key_type = k
-                        .get_type()
-                        .name()
-                        .map(|n| n.to_string())
-                        .unwrap_or_else(|_| "unknown".to_string());
-                    FlussError::new_err(format!("Row dict keys must be strings; got {key_type}"))
-                })?;
-
-                if fields.iter().all(|f| f.name() != key_str) {
-                    let expected = fields
-                        .iter()
-                        .map(|f| f.name())
-                        .collect::<Vec<_>>()
-                        .join(", ");
-                    return Err(FlussError::new_err(format!(
-                        "Unknown field '{key_str}'. Expected fields: {expected}"
-                    )));
-                }
-            }
-
-            let mut datums = Vec::with_capacity(fields.len());
-            for field in fields {
-                let value = dict.get_item(field.name())?.ok_or_else(|| {
-                    FlussError::new_err(format!("Missing field: {}", field.name()))
-                })?;
-                datums.push(
-                    python_value_to_datum(&value, field.data_type()).map_err(|e| {
-                        FlussError::new_err(format!("Field '{}': {}", field.name(), e))
-                    })?,
-                );
-            }
-            datums
-        }
-
-        RowInput::List(list) => process_sequence_to_datums(list.iter(), list.len(), fields)?,
-
-        RowInput::Tuple(tuple) => process_sequence_to_datums(tuple.iter(), tuple.len(), fields)?,
-    };
+    let all_indices: Vec<usize> = (0..table_info.row_type().fields().len()).collect();
+    python_to_sparse_generic_row(row, table_info, &all_indices)
+}
 
-    Ok(fcore::row::GenericRow { values: datums })
+/// Process a Python sequence (list or tuple) into datums at the target column positions.
+fn process_sequence(
+    seq: &Bound<pyo3::types::PySequence>,
+    target_indices: &[usize],
+    fields: &[fcore::metadata::DataField],
+    datums: &mut [fcore::row::Datum<'static>],
+) -> PyResult<()> {
+    if seq.len()? != target_indices.len() {
+        return Err(FlussError::new_err(format!(
+            "Expected {} elements, got {}",
+            target_indices.len(),
+            seq.len()?
+        )));
+    }
+    for (i, &col_idx) in target_indices.iter().enumerate() {
+        let field = &fields[col_idx];
+        let value = seq.get_item(i)?;
+        datums[col_idx] = python_value_to_datum(&value, field.data_type())
+            .map_err(|e| FlussError::new_err(format!("Field '{}': {}", field.name(), e)))?;
+    }
+    Ok(())
 }
 
-/// Convert Python primary key values (dict/list/tuple) to GenericRow.
-/// Only requires PK columns; non-PK columns are filled with Null.
-/// For dict: keys should be PK column names.
-/// For list/tuple: values should be PK values in PK column order.
-pub fn python_pk_to_generic_row(
+/// Build a full-width GenericRow filling only the specified column
+/// indices from user input; all other columns are set to Null.
+pub fn python_to_sparse_generic_row(
     row: &Bound<PyAny>,
     table_info: &fcore::metadata::TableInfo,
+    target_indices: &[usize],
 ) -> PyResult<fcore::row::GenericRow<'static>> {
-    let schema = table_info.get_schema();
     let row_type = table_info.row_type();
     let fields = row_type.fields();
-    let pk_indexes = schema.primary_key_indexes();
-    let pk_names: Vec<&str> = schema.primary_key_column_names();
-
-    if pk_indexes.is_empty() {
-        return Err(FlussError::new_err(
-            "Table has no primary key; cannot use PK-only row",
-        ));
-    }
+    let target_names: Vec<&str> = target_indices.iter().map(|&i| fields[i].name()).collect();
 
-    // Initialize all datums as Null
     let mut datums: Vec<fcore::row::Datum<'static>> = vec![fcore::row::Datum::Null; fields.len()];
 
-    // Extract with user-friendly error message
     let row_input: RowInput = row.extract().map_err(|_| {
         let type_name = row
             .get_type()
             .name()
             .map(|n| n.to_string())
             .unwrap_or_else(|_| "unknown".to_string());
         FlussError::new_err(format!(
-            "PK row must be a dict, list, or tuple; got {type_name}"
+            "Row must be a dict, list, or tuple; got {type_name}"
         ))
     })?;
 
     match row_input {
         RowInput::Dict(dict) => {
-            // Validate keys are PK columns
             for (k, _) in dict.iter() {
                 let key_str = k.extract::<&str>().map_err(|_| {
                     let key_type = k
                         .get_type()
                         .name()
                         .map(|n| n.to_string())
                         .unwrap_or_else(|_| "unknown".to_string());
-                    FlussError::new_err(format!("PK dict keys must be strings; got {key_type}"))
+                    FlussError::new_err(format!("Dict keys must be strings; got {key_type}"))
                 })?;
-
-                if !pk_names.contains(&key_str) {
+                if !target_names.contains(&key_str) {
                     return Err(FlussError::new_err(format!(
-                        "Unknown PK field '{}'. Expected PK fields: {}",
+                        "Unknown field '{}'. Expected: {}",
                         key_str,
-                        pk_names.join(", ")
+                        target_names.join(", ")
                     )));
                 }
             }
-
-            // Extract PK values
-            for (i, pk_idx) in pk_indexes.iter().enumerate() {
-                let pk_name = pk_names[i];
-                let field: &fcore::metadata::DataField = &fields[*pk_idx];
+            for (i, &col_idx) in target_indices.iter().enumerate() {
+                let name = target_names[i];
+                let field = &fields[col_idx];
                 let value = dict
-                    .get_item(pk_name)?
-                    .ok_or_else(|| FlussError::new_err(format!("Missing PK field: {pk_name}")))?;
-                datums[*pk_idx] = python_value_to_datum(&value, field.data_type())
-                    .map_err(|e| FlussError::new_err(format!("PK field '{pk_name}': {e}")))?;
+                    .get_item(name)?
+                    .ok_or_else(|| FlussError::new_err(format!("Missing field: {name}")))?;
+                datums[col_idx] = python_value_to_datum(&value, field.data_type())
+                    .map_err(|e| FlussError::new_err(format!("Field '{name}': {e}")))?;
             }
         }
 
         RowInput::List(list) => {
-            if list.len() != pk_indexes.len() {
-                return Err(FlussError::new_err(format!(
-                    "PK list must have {} elements (PK columns), got {}",
-                    pk_indexes.len(),
-                    list.len()
-                )));
-            }
-            for (i, pk_idx) in pk_indexes.iter().enumerate() {
-                let field: &fcore::metadata::DataField = &fields[*pk_idx];
-                let value = list.get_item(i)?;
-                datums[*pk_idx] =
-                    python_value_to_datum(&value, field.data_type()).map_err(|e| {
-                        FlussError::new_err(format!("PK field '{}': {}", field.name(), e))
-                    })?;
-            }
+            let seq = list.as_sequence();
+            process_sequence(seq, target_indices, fields, &mut datums)?;
         }
 
         RowInput::Tuple(tuple) => {
-            if tuple.len() != pk_indexes.len() {
-                return Err(FlussError::new_err(format!(
-                    "PK tuple must have {} elements (PK columns), got {}",
-                    pk_indexes.len(),
-                    tuple.len()
-                )));
-            }
-            for (i, pk_idx) in pk_indexes.iter().enumerate() {
-                let field: &fcore::metadata::DataField = &fields[*pk_idx];
-                let value = tuple.get_item(i)?;
-                datums[*pk_idx] =
-                    python_value_to_datum(&value, field.data_type()).map_err(|e| {
-                        FlussError::new_err(format!("PK field '{}': {}", field.name(), e))
-                    })?;
-            }
+            let seq = tuple.as_sequence();
+            process_sequence(seq, target_indices, fields, &mut datums)?;
         }
     }
 
diff --git a/bindings/python/src/upsert.rs b/bindings/python/src/upsert.rs
diff --git a/crates/fluss/src/client/table/mod.rs b/crates/fluss/src/client/table/mod.rs
