update based on comments

polyzos · polyzos · commit e7b1bd0855ec · 2025-09-18T15:25:01.000+03:00
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/PartitionFilter.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/PartitionFilter.java
@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+package org.apache.fluss.client.table.scanner;
+
+import javax.annotation.Nullable;
+
+/**
+ * A simple filter that narrows scans to a specific partition.
+ *
+ * <p>This is an initial, minimal filter abstraction to enable partition pruning for client-side
+ * scans. For now, only equality on a concrete partition name is supported. Row-level filtering is
+ * intentionally not supported here.
+ */
+public final class PartitionFilter {
+
+    @Nullable private final String partitionName;
+
+    private PartitionFilter(@Nullable String partitionName) {
+        this.partitionName = partitionName;
+    }
+
+    /** Creates a filter that limits the scan to the given partition name. */
+    public static PartitionFilter ofPartitionName(String partitionName) {
+        return new PartitionFilter(partitionName);
+    }
+
+    /** Returns the partition name to filter by, or null if none. */
+    @Nullable
+    public String getPartitionName() {
+        return partitionName;
+    }
+}
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/Scan.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/Scan.java
@@ -51,6 +51,14 @@ public interface Scan {
      */
     Scan project(List<String> projectedColumnNames);
 
+    /**
+     * Returns a new scan from this refined by a partition filter.
+     *
+     * <p>Currently, only equality predicates on partition columns are supported. Non-partitioned
+     * tables do not accept filters.
+     */
+    Scan filter(PartitionFilter partitionFilter);
+
     /**
      * Returns a new scan from this that will read the given limited row number.
      *
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/TableScan.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/TableScan.java
@@ -21,6 +21,7 @@
 import org.apache.fluss.client.admin.Admin;
 import org.apache.fluss.client.metadata.KvSnapshotMetadata;
 import org.apache.fluss.client.table.scanner.batch.BatchScanner;
+import org.apache.fluss.client.table.scanner.batch.DefaultBatchScanner;
 import org.apache.fluss.client.table.scanner.batch.KvSnapshotBatchScanner;
 import org.apache.fluss.client.table.scanner.batch.LimitBatchScanner;
 import org.apache.fluss.client.table.scanner.log.LogScanner;
@@ -43,27 +44,31 @@ public class TableScan implements Scan {
 
     /** The projected fields to do projection. No projection if is null. */
     @Nullable private final int[] projectedColumns;
+    /** Partition filter for partition pruning; null if none. */
+    @Nullable private final PartitionFilter partitionFilter;
     /** The limited row number to read. No limit if is null. */
     @Nullable private final Integer limit;
 
     public TableScan(FlussConnection conn, TableInfo tableInfo) {
-        this(conn, tableInfo, null, null);
+        this(conn, tableInfo, null, null, null);
     }
 
     private TableScan(
             FlussConnection conn,
             TableInfo tableInfo,
             @Nullable int[] projectedColumns,
+            @Nullable PartitionFilter partitionFilter,
             @Nullable Integer limit) {
         this.conn = conn;
         this.tableInfo = tableInfo;
         this.projectedColumns = projectedColumns;
+        this.partitionFilter = partitionFilter;
         this.limit = limit;
     }
 
     @Override
     public Scan project(@Nullable int[] projectedColumns) {
-        return new TableScan(conn, tableInfo, projectedColumns, limit);
+        return new TableScan(conn, tableInfo, projectedColumns, partitionFilter, limit);
     }
 
     @Override
@@ -78,12 +83,21 @@ public Scan project(List<String> projectedColumnNames) {
             }
             columnIndexes[i] = index;
         }
-        return new TableScan(conn, tableInfo, columnIndexes, limit);
+        return new TableScan(conn, tableInfo, columnIndexes, partitionFilter, limit);
+    }
+
+    @Override
+    public Scan filter(PartitionFilter partitionFilter) {
+        if (!tableInfo.isPartitioned()) {
+            throw new UnsupportedOperationException(
+                    "Partition filter is only supported for partitioned tables.");
+        }
+        return new TableScan(conn, tableInfo, projectedColumns, partitionFilter, limit);
     }
 
     @Override
     public Scan limit(int rowNumber) {
-        return new TableScan(conn, tableInfo, projectedColumns, rowNumber);
+        return new TableScan(conn, tableInfo, projectedColumns, partitionFilter, rowNumber);
     }
 
     @Override
@@ -103,8 +117,13 @@ public LogScanner createLogScanner() {
     @Override
     public BatchScanner createBatchScanner(TableBucket tableBucket) {
         if (limit == null) {
-            throw new UnsupportedOperationException(
-                    "Currently, BatchScanner is only available when limit is set.");
+            if (!tableInfo.hasPrimaryKey()) {
+                throw new UnsupportedOperationException(
+                        "Full scan BatchScanner is only supported for primary key tables.");
+            }
+            // Full scan semantics: one-shot snapshot of current values across all buckets.
+            return new DefaultBatchScanner(
+                    tableInfo, conn.getMetadataUpdater(), projectedColumns, partitionFilter);
         }
         return new LimitBatchScanner(
                 tableInfo, tableBucket, conn.getMetadataUpdater(), projectedColumns, limit);
diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/DefaultBatchScanner.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/DefaultBatchScanner.java
@@ -0,0 +1,217 @@
+/*
+ * 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.
+ */
+
+package org.apache.fluss.client.table.scanner.batch;
+
+import org.apache.fluss.client.metadata.MetadataUpdater;
+import org.apache.fluss.client.table.scanner.PartitionFilter;
+import org.apache.fluss.exception.LeaderNotAvailableException;
+import org.apache.fluss.exception.TableNotPartitionedException;
+import org.apache.fluss.metadata.PhysicalTablePath;
+import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.metadata.TableInfo;
+import org.apache.fluss.metadata.TablePath;
+import org.apache.fluss.record.DefaultValueRecordBatch;
+import org.apache.fluss.record.ValueRecord;
+import org.apache.fluss.record.ValueRecordReadContext;
+import org.apache.fluss.row.GenericRow;
+import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.row.ProjectedRow;
+import org.apache.fluss.row.decode.RowDecoder;
+import org.apache.fluss.row.encode.ValueDecoder;
+import org.apache.fluss.rpc.gateway.TabletServerGateway;
+import org.apache.fluss.rpc.messages.FullScanRequest;
+import org.apache.fluss.rpc.messages.FullScanResponse;
+import org.apache.fluss.rpc.protocol.Errors;
+import org.apache.fluss.types.DataType;
+import org.apache.fluss.types.RowType;
+import org.apache.fluss.utils.CloseableIterator;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * Default implementation of {@link BatchScanner} that performs a full scan against tablet servers.
+ *
+ * <p>This scanner issues FULL_SCAN RPCs to the leaders of all buckets and aggregates the results.
+ * It returns all current values at a point in time for primary-key tables. The first call to
+ * {@link #pollBatch(Duration)} returns the complete snapshot; subsequent calls return {@code null}.
+ *
+ * <p>Note: For partitioned tables, callers may provide a {@link PartitionFilter} with a partition
+ * name to restrict the scan to a single partition.
+ */
+public class DefaultBatchScanner implements BatchScanner {
+
+    private final TableInfo tableInfo;
+    private final MetadataUpdater metadataUpdater;
+    @Nullable private final int[] projectedFields;
+    @Nullable private final PartitionFilter partitionFilter;
+
+    private final InternalRow.FieldGetter[] fieldGetters;
+    private final ValueDecoder kvValueDecoder;
+
+    private boolean endOfInput = false;
+
+    public DefaultBatchScanner(
+            TableInfo tableInfo,
+            MetadataUpdater metadataUpdater,
+            @Nullable int[] projectedFields,
+            @Nullable PartitionFilter partitionFilter) {
+        this.tableInfo = Objects.requireNonNull(tableInfo, "tableInfo");
+        this.metadataUpdater = Objects.requireNonNull(metadataUpdater, "metadataUpdater");
+        this.projectedFields = projectedFields;
+        this.partitionFilter = partitionFilter;
+
+        RowType rowType = tableInfo.getRowType();
+        this.fieldGetters = new InternalRow.FieldGetter[rowType.getFieldCount()];
+        for (int i = 0; i < rowType.getFieldCount(); i++) {
+            this.fieldGetters[i] = InternalRow.createFieldGetter(rowType.getTypeAt(i), i);
+        }
+        this.kvValueDecoder =
+                new ValueDecoder(
+                        RowDecoder.create(
+                                tableInfo.getTableConfig().getKvFormat(),
+                                rowType.getChildren().toArray(new DataType[0])));
+    }
+
+    @Nullable
+    @Override
+    public CloseableIterator<InternalRow> pollBatch(Duration timeout) throws IOException {
+        if (endOfInput) {
+            return null;
+        }
+        try {
+            List<CompletableFuture<FullScanResponse>> futures = issueFullScanRequests();
+            // wait for all responses or timeout for this poll
+            CompletableFuture<Void> all =
+                    CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]));
+            all.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
+            List<InternalRow> rows = decodeFullScanResponses(futures);
+            endOfInput = true;
+            return CloseableIterator.wrap(rows.iterator());
+        } catch (TimeoutException e) {
+            // try again in next poll
+            return CloseableIterator.emptyIterator();
+        } catch (Exception e) {
+            throw new IOException(e);
+        }
+    }
+
+    private List<CompletableFuture<FullScanResponse>> issueFullScanRequests() {
+        Long partitionId = null;
+        if (tableInfo.isPartitioned()) {
+            if (partitionFilter == null || partitionFilter.getPartitionName() == null) {
+                throw new TableNotPartitionedException(
+                        "Partition filter is required for partitioned table full scan.");
+            }
+            TablePath tablePath = tableInfo.getTablePath();
+            PhysicalTablePath physicalTablePath =
+                    PhysicalTablePath.of(tablePath, partitionFilter.getPartitionName());
+            metadataUpdater.checkAndUpdatePartitionMetadata(physicalTablePath);
+            partitionId =
+                    metadataUpdater
+                            .getPartitionId(physicalTablePath)
+                            .orElseThrow(
+                                    () ->
+                                            new IllegalStateException(
+                                                    "Partition id not found for "
+                                                            + partitionFilter.getPartitionName()));
+        }
+
+        long tableId = tableInfo.getTableId();
+        int numBuckets = tableInfo.getNumBuckets();
+
+        // collect leaders for all buckets
+        HashSet<Integer> leaderServers = new HashSet<>();
+        for (int bucketId = 0; bucketId < numBuckets; bucketId++) {
+            TableBucket tableBucket = new TableBucket(tableId, partitionId, bucketId);
+            metadataUpdater.checkAndUpdateMetadata(tableInfo.getTablePath(), tableBucket);
+            int leader = metadataUpdater.leaderFor(tableBucket);
+            leaderServers.add(leader);
+        }
+
+        List<CompletableFuture<FullScanResponse>> responseFutures = new ArrayList<>();
+        for (int leader : leaderServers) {
+            TabletServerGateway gateway = metadataUpdater.newTabletServerClientForNode(leader);
+            if (gateway == null) {
+                throw new LeaderNotAvailableException(
+                        "Server " + leader + " is not found in metadata cache.");
+            }
+            FullScanRequest request = new FullScanRequest().setTableId(tableId);
+            if (partitionId != null) {
+                request.setPartitionId(partitionId);
+            }
+            // Future-proof: optionally pass bucket list when supported by server
+            // request.setBucketId(Collections.singletonList(bucketId));
+            responseFutures.add(gateway.fullScan(request));
+        }
+        return responseFutures;
+    }
+
+    private List<InternalRow> decodeFullScanResponses(
+            List<CompletableFuture<FullScanResponse>> responseFutures) {
+        List<InternalRow> out = new ArrayList<>();
+        for (CompletableFuture<FullScanResponse> responseFuture : responseFutures) {
+            FullScanResponse response = responseFuture.join();
+            if (response.hasErrorCode() && response.getErrorCode() != Errors.NONE.code()) {
+                Errors err = Errors.forCode(response.getErrorCode());
+                throw err.exception(
+                        response.hasErrorMessage() ? response.getErrorMessage() : err.message());
+            }
+            if (response.hasRecords()) {
+                ByteBuffer buffer = ByteBuffer.wrap(response.getRecords());
+                DefaultValueRecordBatch values = DefaultValueRecordBatch.pointToByteBuffer(buffer);
+                ValueRecordReadContext context =
+                        new ValueRecordReadContext(kvValueDecoder.getRowDecoder());
+                for (ValueRecord record : values.records(context)) {
+                    out.add(maybeProject(record.getRow()));
+                }
+            }
+        }
+        return out;
+    }
+
+    private InternalRow maybeProject(InternalRow originRow) {
+        // deep copy and project if requested to avoid referencing released buffers
+        GenericRow newRow = new GenericRow(fieldGetters.length);
+        for (int i = 0; i < fieldGetters.length; i++) {
+            newRow.setField(i, fieldGetters[i].getFieldOrNull(originRow));
+        }
+        if (projectedFields != null) {
+            ProjectedRow projectedRow = ProjectedRow.from(projectedFields);
+            projectedRow.replaceRow(newRow);
+            return projectedRow;
+        } else {
+            return newRow;
+        }
+    }
+
+    @Override
+    public void close() {
+        // no-op
+    }
+}
