[Draft][Discussion] Support implicit partition

fluss-client/src/main/java/org/apache/fluss/client/lookup/PrefixKeyLookuper.java

@@ -19,7 +19,6 @@
 
 import org.apache.fluss.bucketing.BucketingFunction;
 import org.apache.fluss.client.metadata.MetadataUpdater;
-import org.apache.fluss.client.table.getter.PartitionGetter;
 import org.apache.fluss.exception.PartitionNotExistException;
 import org.apache.fluss.metadata.DataLakeFormat;
 import org.apache.fluss.metadata.SchemaGetter;
@@ -28,6 +27,7 @@
 import org.apache.fluss.row.InternalRow;
 import org.apache.fluss.row.encode.KeyEncoder;
 import org.apache.fluss.types.RowType;
+import org.apache.fluss.utils.PartitionComputer;
 
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
@@ -60,7 +60,7 @@ class PrefixKeyLookuper extends AbstractLookuper implements Lookuper {
     /**
      * a getter to extract partition from prefix lookup key row, null when it's not a partitioned.
      */
-    private @Nullable final PartitionGetter partitionGetter;
+    private @Nullable final PartitionComputer partitionComputer;
 
     public PrefixKeyLookuper(
             TableInfo tableInfo,
@@ -95,10 +95,8 @@ public PrefixKeyLookuper(
                                 lookupRowType, tableInfo.getBucketKeys(), lakeFormat);
 
         this.bucketingFunction = BucketingFunction.of(lakeFormat);
-        this.partitionGetter =
-                tableInfo.isPartitioned()
-                        ? new PartitionGetter(lookupRowType, tableInfo.getPartitionKeys())
-                        : null;
+        this.partitionComputer =
+                tableInfo.isPartitioned() ? new PartitionComputer(tableInfo, lookupRowType) : null;
     }
 
     private void validatePrefixLookup(TableInfo tableInfo, List<String> lookupColumns) {
@@ -126,7 +124,7 @@ private void validatePrefixLookup(TableInfo tableInfo, List<String> lookupColumn
 
         // verify the lookup columns must contain all partition fields if this is partitioned table
         if (tableInfo.isPartitioned()) {
-            List<String> partitionKeys = tableInfo.getPartitionKeys();
+            List<String> partitionKeys = tableInfo.getPartitionInputColumns();
             Set<String> lookupColumnsSet = new HashSet<>(lookupColumns);
             if (!lookupColumnsSet.containsAll(partitionKeys)) {
                 throw new IllegalArgumentException(
@@ -138,9 +136,9 @@ private void validatePrefixLookup(TableInfo tableInfo, List<String> lookupColumn
         }
 
         // verify the lookup columns must contain all bucket keys **in order**
-        List<String> physicalLookupColumns = new ArrayList<>(lookupColumns);
-        physicalLookupColumns.removeAll(tableInfo.getPartitionKeys());
-        if (!physicalLookupColumns.equals(bucketKeys)) {
+        List<String> bucketLookupColumns =
+                removePartitionOnlyInputColumns(tableInfo, lookupColumns);
+        if (!bucketLookupColumns.equals(bucketKeys)) {
             throw new IllegalArgumentException(
                     String.format(
                             "Can not perform prefix lookup on table '%s', "
@@ -158,6 +156,21 @@ private void validatePrefixLookup(TableInfo tableInfo, List<String> lookupColumn
         }
     }
 
+    private List<String> removePartitionOnlyInputColumns(
+            TableInfo tableInfo, List<String> lookupColumns) {
+        Set<String> bucketKeySet = new HashSet<>(tableInfo.getBucketKeys());
+        Set<String> partitionInputColumnSet = new HashSet<>(tableInfo.getPartitionInputColumns());
+        List<String> bucketLookupColumns = new ArrayList<>();
+        for (String lookupColumn : lookupColumns) {
+            if (partitionInputColumnSet.contains(lookupColumn)
+                    && !bucketKeySet.contains(lookupColumn)) {
+                continue;
+            }
+            bucketLookupColumns.add(lookupColumn);
+        }
+        return bucketLookupColumns;
+    }
+
     @Override
     public CompletableFuture<LookupResult> lookup(InternalRow prefixKey) {
         byte[] prefixKeyBytes = prefixKeyEncoder.encodeKey(prefixKey);
@@ -168,12 +181,12 @@ public CompletableFuture<LookupResult> lookup(InternalRow prefixKey) {
         int bucketId = bucketingFunction.bucketing(bucketKeyBytes, numBuckets);
 
         Long partitionId = null;
-        if (partitionGetter != null) {
+        if (partitionComputer != null) {
             try {
                 partitionId =
                         getPartitionId(
                                 prefixKey,
-                                partitionGetter,
+                                partitionComputer,
                                 tableInfo.getTablePath(),
                                 metadataUpdater);
             } catch (PartitionNotExistException e) {

fluss-client/src/main/java/org/apache/fluss/client/lookup/PrimaryKeyLookuper.java

@@ -19,7 +19,6 @@
 
 import org.apache.fluss.bucketing.BucketingFunction;
 import org.apache.fluss.client.metadata.MetadataUpdater;
-import org.apache.fluss.client.table.getter.PartitionGetter;
 import org.apache.fluss.exception.PartitionNotExistException;
 import org.apache.fluss.metadata.DataLakeFormat;
 import org.apache.fluss.metadata.SchemaGetter;
@@ -28,6 +27,7 @@
 import org.apache.fluss.row.InternalRow;
 import org.apache.fluss.row.encode.KeyEncoder;
 import org.apache.fluss.types.RowType;
+import org.apache.fluss.utils.PartitionComputer;
 
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
@@ -55,7 +55,7 @@ class PrimaryKeyLookuper extends AbstractLookuper implements Lookuper {
     private final boolean insertIfNotExists;
 
     /** a getter to extract partition from lookup key row, null when it's not a partitioned. */
-    private @Nullable final PartitionGetter partitionGetter;
+    private @Nullable final PartitionComputer partitionComputer;
 
     public PrimaryKeyLookuper(
             TableInfo tableInfo,
@@ -90,10 +90,8 @@ public PrimaryKeyLookuper(
 
         this.bucketingFunction = BucketingFunction.of(lakeFormat);
 
-        this.partitionGetter =
-                tableInfo.isPartitioned()
-                        ? new PartitionGetter(lookupRowType, tableInfo.getPartitionKeys())
-                        : null;
+        this.partitionComputer =
+                tableInfo.isPartitioned() ? new PartitionComputer(tableInfo, lookupRowType) : null;
     }
 
     @Override
@@ -106,12 +104,12 @@ public CompletableFuture<LookupResult> lookup(InternalRow lookupKey) {
                         ? pkBytes
                         : bucketKeyEncoder.encodeKey(lookupKey);
         Long partitionId = null;
-        if (partitionGetter != null) {
+        if (partitionComputer != null) {
             try {
                 partitionId =
                         getPartitionId(
                                 lookupKey,
-                                partitionGetter,
+                                partitionComputer,
                                 tableInfo.getTablePath(),
                                 metadataUpdater);
             } catch (PartitionNotExistException e) {

fluss-client/src/main/java/org/apache/fluss/client/metadata/MetadataUpdater.java

@@ -182,7 +182,10 @@ public void checkAndUpdateTableMetadata(Set<TablePath> tablePaths) {
     public boolean checkAndUpdatePartitionMetadata(PhysicalTablePath physicalTablePath)
             throws PartitionNotExistException {
         if (!cluster.getPartitionId(physicalTablePath).isPresent()) {
-            updateMetadata(null, Collections.singleton(physicalTablePath), null);
+            updateMetadata(
+                    Collections.singleton(physicalTablePath.getTablePath()),
+                    Collections.singleton(physicalTablePath),
+                    null);
         }
         return cluster.getPartitionId(physicalTablePath).isPresent();
     }

fluss-client/src/main/java/org/apache/fluss/client/table/writer/AbstractTableWriter.java

@@ -17,7 +17,6 @@
 
 package org.apache.fluss.client.table.writer;
 
-import org.apache.fluss.client.table.getter.PartitionGetter;
 import org.apache.fluss.client.write.WriteRecord;
 import org.apache.fluss.client.write.WriterClient;
 import org.apache.fluss.config.ConfigOptions;
@@ -26,6 +25,7 @@
 import org.apache.fluss.metadata.TableInfo;
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.utils.PartitionComputer;
 
 import javax.annotation.Nullable;
 
@@ -38,16 +38,16 @@ public abstract class AbstractTableWriter implements TableWriter {
     protected final TablePath tablePath;
     protected final WriterClient writerClient;
     protected final int fieldCount;
-    private final @Nullable PartitionGetter partitionFieldGetter;
+    private final @Nullable PartitionComputer partitionComputer;
 
     protected AbstractTableWriter(
             TablePath tablePath, TableInfo tableInfo, WriterClient writerClient) {
         this.tablePath = tablePath;
         this.writerClient = writerClient;
         this.fieldCount = tableInfo.getRowType().getFieldCount();
-        this.partitionFieldGetter =
+        this.partitionComputer =
                 tableInfo.isPartitioned()
-                        ? new PartitionGetter(tableInfo.getRowType(), tableInfo.getPartitionKeys())
+                        ? new PartitionComputer(tableInfo, tableInfo.getRowType())
                         : null;
     }
 
@@ -108,11 +108,11 @@ protected interface ResultFactory<T> {
 
     protected PhysicalTablePath getPhysicalPath(InternalRow row) {
         // not partitioned table, return the original physical path
-        if (partitionFieldGetter == null) {
+        if (partitionComputer == null) {
             return PhysicalTablePath.of(tablePath);
         } else {
             // partitioned table, extract partition from the row
-            String partition = partitionFieldGetter.getPartition(row);
+            String partition = partitionComputer.getPartition(row);
             return PhysicalTablePath.of(tablePath, partition);
         }
     }

fluss-client/src/main/java/org/apache/fluss/client/table/writer/TypedUpsertWriterImpl.java

@@ -34,7 +34,6 @@
 class TypedUpsertWriterImpl<T> implements TypedUpsertWriter<T> {
 
     private final UpsertWriter delegate;
-    private final TableInfo tableInfo;
     private final RowType tableSchema;
     @Nullable private final int[] targetColumns;
 
@@ -48,12 +47,11 @@ class TypedUpsertWriterImpl<T> implements TypedUpsertWriter<T> {
     TypedUpsertWriterImpl(
             UpsertWriter delegate, Class<T> pojoClass, TableInfo tableInfo, int[] targetColumns) {
         this.delegate = delegate;
-        this.tableInfo = tableInfo;
         this.tableSchema = tableInfo.getRowType();
         this.targetColumns = targetColumns;
 
         // Precompute projections
-        this.pkProjection = this.tableSchema.project(tableInfo.getPhysicalPrimaryKeys());
+        this.pkProjection = this.tableSchema.project(tableInfo.getPrimaryKeys());
         this.targetProjection =
                 (targetColumns == null) ? null : this.tableSchema.project(targetColumns);
 
@@ -92,7 +90,10 @@ public CompletableFuture<DeleteResult> delete(T record) {
     private InternalRow convertPojo(T pojo, boolean forDelete) {
         final RowType projection;
         final PojoToRowConverter<T> converter;
-        if (forDelete) {
+        if (forDelete && pkProjection.getFieldCount() == tableSchema.getFieldCount()) {
+            projection = tableSchema;
+            converter = pojoToRowConverter;
+        } else if (forDelete) {
             projection = pkProjection;
             converter = pkConverter;
         } else if (targetProjection != null && targetConverter != null) {
@@ -104,22 +105,12 @@ private InternalRow convertPojo(T pojo, boolean forDelete) {
         }
 
         GenericRow projected = converter.toRow(pojo);
-        if (projection == tableSchema) {
+        if (projection == tableSchema || forDelete) {
             return projected;
         }
         // expand projected row to full row if needed
         GenericRow full = new GenericRow(tableSchema.getFieldCount());
-        if (forDelete) {
-            // set PK fields, others null
-            for (String pk : tableInfo.getPhysicalPrimaryKeys()) {
-                int projIndex = projection.getFieldIndex(pk);
-
-                // TODO: this can be optimized by pre-computing
-                // the index mapping in the constructor?
-                int fullIndex = tableSchema.getFieldIndex(pk);
-                full.setField(fullIndex, projected.getField(projIndex));
-            }
-        } else if (targetColumns != null) {
+        if (targetColumns != null) {
             for (int i = 0; i < projection.getFieldCount(); i++) {
                 String name = projection.getFieldNames().get(i);
                 int fullIdx = tableSchema.getFieldIndex(name);

fluss-client/src/main/java/org/apache/fluss/client/table/writer/UpsertWriter.java

@@ -40,6 +40,8 @@ public interface UpsertWriter extends TableWriter {
 
     /**
      * Delete a certain record from the Fluss table. The input must contain the primary key fields.
+     * When the full table row and primary-key row have the same field count, the input is treated
+     * as a full table row in table schema order.
      *
      * @param record the record to delete.
      * @return A {@link CompletableFuture} that always delete result when complete normally.

fluss-client/src/main/java/org/apache/fluss/client/table/writer/UpsertWriterImpl.java

@@ -21,6 +21,7 @@
 import org.apache.fluss.client.write.WriteRecord;
 import org.apache.fluss.client.write.WriterClient;
 import org.apache.fluss.metadata.KvFormat;
+import org.apache.fluss.metadata.PhysicalTablePath;
 import org.apache.fluss.metadata.TableInfo;
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.row.BinaryRow;
@@ -32,6 +33,7 @@
 import org.apache.fluss.row.indexed.IndexedRow;
 import org.apache.fluss.rpc.protocol.MergeMode;
 import org.apache.fluss.types.RowType;
+import org.apache.fluss.utils.PartitionComputer;
 
 import javax.annotation.Nullable;
 
@@ -45,6 +47,10 @@ class UpsertWriterImpl extends AbstractTableWriter implements UpsertWriter {
     private final TableInfo tableInfo;
     private final KeyEncoder primaryKeyEncoder;
     private final @Nullable int[] targetColumns;
+    private final RowType primaryKeyRowType;
+    private final KeyEncoder primaryKeyDeleteEncoder;
+    private final KeyEncoder bucketKeyDeleteEncoder;
+    private final @Nullable PartitionComputer deletePartitionComputer;
 
     // same to primaryKeyEncoder if the bucket key is the same to the primary key
     private final KeyEncoder bucketKeyEncoder;
@@ -94,6 +100,24 @@ class UpsertWriterImpl extends AbstractTableWriter implements UpsertWriter {
                         tableInfo.getTableConfig(),
                         tableInfo.isDefaultBucketKey(),
                         primaryKeyEncoder);
+        this.primaryKeyRowType = rowType.project(tableInfo.getPrimaryKeys());
+        this.primaryKeyDeleteEncoder =
+                KeyEncoder.ofPrimaryKeyEncoder(
+                        primaryKeyRowType,
+                        tableInfo.getPhysicalPrimaryKeys(),
+                        tableInfo.getTableConfig(),
+                        tableInfo.isDefaultBucketKey());
+        this.bucketKeyDeleteEncoder =
+                KeyEncoder.ofBucketKeyEncoder(
+                        primaryKeyRowType,
+                        tableInfo.getBucketKeys(),
+                        tableInfo.getTableConfig(),
+                        tableInfo.isDefaultBucketKey(),
+                        primaryKeyDeleteEncoder);
+        this.deletePartitionComputer =
+                tableInfo.isPartitioned()
+                        ? new PartitionComputer(tableInfo, primaryKeyRowType)
+                        : null;
 
         this.kvFormat = tableInfo.getTableConfig().getKvFormat();
         this.writeFormat = WriteFormat.fromKvFormat(this.kvFormat);
@@ -201,7 +225,25 @@ public CompletableFuture<UpsertResult> upsert(InternalRow row) {
      */
     @Override
     public CompletableFuture<DeleteResult> delete(InternalRow row) {
-        checkFieldCount(row);
+        // Prefer full-row interpretation when full table row and primary-key row have the same
+        // field count. This preserves table-schema ordering for all-column primary-key tables.
+        if (row.getFieldCount() == fieldCount) {
+            return deleteFullRow(row);
+        }
+        if (row.getFieldCount() == primaryKeyRowType.getFieldCount()) {
+            return deletePrimaryKeyRow(row);
+        }
+        throw new IllegalArgumentException(
+                "The field count of the row does not match the table schema or primary key schema. "
+                        + "Expected full table row: "
+                        + fieldCount
+                        + ", expected primary key row: "
+                        + primaryKeyRowType.getFieldCount()
+                        + ", Actual: "
+                        + row.getFieldCount());
+    }
+
+    private CompletableFuture<DeleteResult> deleteFullRow(InternalRow row) {
         byte[] key = primaryKeyEncoder.encodeKey(row);
         byte[] bucketKey =
                 bucketKeyEncoder == primaryKeyEncoder ? key : bucketKeyEncoder.encodeKey(row);
@@ -217,6 +259,31 @@ public CompletableFuture<DeleteResult> delete(InternalRow row) {
         return sendWithResult(record, DeleteResult::new);
     }
 
+    private CompletableFuture<DeleteResult> deletePrimaryKeyRow(InternalRow row) {
+        byte[] key = primaryKeyDeleteEncoder.encodeKey(row);
+        byte[] bucketKey =
+                bucketKeyDeleteEncoder == primaryKeyDeleteEncoder
+                        ? key
+                        : bucketKeyDeleteEncoder.encodeKey(row);
+        WriteRecord record =
+                WriteRecord.forDelete(
+                        tableInfo,
+                        getDeletePhysicalPath(row),
+                        key,
+                        bucketKey,
+                        writeFormat,
+                        targetColumns,
+                        mergeMode);
+        return sendWithResult(record, DeleteResult::new);
+    }
+
+    private PhysicalTablePath getDeletePhysicalPath(InternalRow row) {
+        if (deletePartitionComputer == null) {
+            return PhysicalTablePath.of(tablePath);
+        }
+        return PhysicalTablePath.of(tablePath, deletePartitionComputer.getPartition(row));
+    }
+
     private BinaryRow encodeRow(InternalRow row) {
         if (kvFormat == KvFormat.INDEXED && row instanceof IndexedRow) {
             return (IndexedRow) row;