Introduce UpsertWriter

crates/fluss/Cargo.toml

@@ -33,6 +33,7 @@ integration_tests = []
 [dependencies]
 arrow = { workspace = true }
 arrow-schema = "57.0.0"
+bitvec = "1"
 byteorder = "1.5"
 futures = "0.3"
 clap = { workspace = true }

crates/fluss/src/client/table/log_fetch_buffer.rs

@@ -651,14 +651,14 @@ mod tests {
     use crate::compression::{
         ArrowCompressionInfo, ArrowCompressionType, DEFAULT_NON_ZSTD_COMPRESSION_LEVEL,
     };
-    use crate::metadata::{DataField, DataTypes, TablePath};
+    use crate::metadata::{DataField, DataTypes, RowType, TablePath};
     use crate::record::{MemoryLogRecordsArrowBuilder, ReadContext, to_arrow_schema};
     use crate::row::GenericRow;
     use std::sync::Arc;
     use std::time::Duration;
 
     fn test_read_context() -> ReadContext {
-        let row_type = DataTypes::row(vec![DataField::new(
+        let row_type = RowType::new(vec![DataField::new(
             "id".to_string(),
             DataTypes::int(),
             None,
@@ -714,7 +714,7 @@ mod tests {
 
     #[test]
     fn default_completed_fetch_reads_records() -> Result<()> {
-        let row_type = DataTypes::row(vec![
+        let row_type = RowType::new(vec![
             DataField::new("id".to_string(), DataTypes::int(), None),
             DataField::new("name".to_string(), DataTypes::string(), None),
         ]);

crates/fluss/src/client/table/lookup.rs

@@ -22,7 +22,7 @@ use crate::error::{Error, Result};
 use crate::metadata::{RowType, TableBucket, TableInfo};
 use crate::row::InternalRow;
 use crate::row::compacted::CompactedRow;
-use crate::row::encode::KeyEncoder;
+use crate::row::encode::{KeyEncoder, KeyEncoderFactory};
 use crate::rpc::ApiError;
 use crate::rpc::message::LookupRequest;
 use std::sync::Arc;
@@ -130,8 +130,11 @@ impl<'a> TableLookup<'a> {
 
         // Create key encoder for the primary key fields
         let pk_fields = self.table_info.get_physical_primary_keys().to_vec();
-        let key_encoder =
-            <dyn KeyEncoder>::of(self.table_info.row_type(), pk_fields, data_lake_format)?;
+        let key_encoder = KeyEncoderFactory::of(
+            self.table_info.row_type(),
+            pk_fields.as_slice(),
+            &data_lake_format,
+        )?;
 
         Ok(Lookuper {
             conn: self.conn,

crates/fluss/src/client/table/mod.rs

@@ -27,13 +27,17 @@ mod append;
 mod lookup;
 
 mod log_fetch_buffer;
+mod partition_getter;
 mod remote_log;
 mod scanner;
+mod upsert;
 mod writer;
 
+use crate::client::table::upsert::TableUpsert;
 pub use append::{AppendWriter, TableAppend};
 pub use lookup::{LookupResult, Lookuper, TableLookup};
 pub use scanner::{LogScanner, RecordBatchLogScanner, TableScan};
+pub use writer::{TableWriter, UpsertWriter};
 
 #[allow(dead_code)]
 pub struct FlussTable<'a> {
@@ -119,6 +123,20 @@ impl<'a> FlussTable<'a> {
             self.metadata.clone(),
         ))
     }
+
+    pub fn new_upsert(&self) -> Result<TableUpsert> {
+        if !self.has_primary_key {
+            return Err(Error::UnsupportedOperation {
+                message: "Upsert is only supported for primary key tables".to_string(),
+            });
+        }
+
+        Ok(TableUpsert::new(
+            self.table_path.clone(),
+            self.table_info.clone(),
+            self.conn.get_or_create_writer_client()?,
+        ))
+    }
 }
 
 impl<'a> Drop for FlussTable<'a> {

crates/fluss/src/client/table/partition_getter.rs

@@ -0,0 +1,56 @@
+// 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.
+
+use crate::error::Error::IllegalArgument;
+use crate::error::Result;
+use crate::metadata::{DataType, RowType};
+use crate::row::field_getter::FieldGetter;
+
+#[allow(dead_code)]
+pub struct PartitionGetter<'a> {
+    partitions: Vec<(&'a String, &'a DataType, FieldGetter)>,
+}
+
+#[allow(dead_code)]
+impl<'a> PartitionGetter<'a> {
+    pub fn new(row_type: &'a RowType, partition_keys: &'a Vec<String>) -> Result<Self> {
+        let mut partitions = Vec::with_capacity(partition_keys.len());
+
+        for partition_key in partition_keys {
+            if let Some(partition_col_index) = row_type.get_field_index(partition_key.as_str()) {
+                let data_type = &row_type
+                    .fields()
+                    .get(partition_col_index)
+                    .unwrap()
+                    .data_type;
+                let field_getter = FieldGetter::create(data_type, partition_col_index);
+
+                partitions.push((partition_key, data_type, field_getter));
+            } else {
+                return Err(IllegalArgument {
+                    message: format!(
+                        "The partition column {partition_key} is not in the row {row_type}."
+                    ),
+                });
+            };
+        }
+
+        Ok(Self { partitions })
+    }
+
+    // TODO Implement get partition
+}