Bulk Export Query Splitter

java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/model/BulkExportLeafPartitionQuery.java → java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/core/model/BulkExportLeafPartitionQuery.java

@@ -14,13 +14,17 @@
  * limitations under the License.
  */
 
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
+import sleeper.core.partition.Partition;
+import sleeper.core.properties.table.TableProperties;
 import sleeper.core.range.Region;
 
 import java.util.List;
 import java.util.Objects;
 
+import static sleeper.core.properties.table.TableProperty.TABLE_ID;
+
 /**
  * An export query for a leaf partition. The query contains information about
  * which files should be read.
@@ -65,14 +69,31 @@ public BulkExportLeafPartitionQuery validate() {
     }
 
     /**
-     * Builder class for BulkExportLeafPartitionQuery.
+     * Creates a builder for this class.
      *
-     * @return a Builder object
+     * @return the builder
      */
     public static Builder builder() {
         return new Builder();
     }
 
+    /**
+     * Creates a builder for this class, pre-initialised for a given partition.
+     *
+     * @param  parentQuery     the parent query
+     * @param  tableProperties the Sleeper table properties
+     * @param  partition       the partition
+     * @return                 the builder
+     */
+    public static Builder forPartition(BulkExportQuery parentQuery, TableProperties tableProperties, Partition partition) {
+        return builder()
+                .exportId(parentQuery.getExportId())
+                .tableId(tableProperties.get(TABLE_ID))
+                .leafPartitionId(partition.getId())
+                .regions(List.of(partition.getRegion()))
+                .partitionRegion(partition.getRegion());
+    }
+
     /**
      * Gets the table Id.
      *
@@ -161,7 +182,7 @@ public int hashCode() {
 
     @Override
     public String toString() {
-        return "LeafPartitionQuery{" +
+        return "BulkExportLeafPartitionQuery{" +
                 "tableId='" + tableId + '\'' +
                 ", exportId='" + exportId + '\'' +
                 ", subExportId='" + subExportId + '\'' +
@@ -180,7 +201,7 @@ private static <T> T requireNonNull(T obj, Builder builder, String message) {
     }
 
     /**
-     * Builder for the BulkExportLeafPartitionQuery model.
+     * Builder for this class.
      */
     public static final class Builder {
         private String tableId;
@@ -195,11 +216,10 @@ private Builder() {
         }
 
         /**
-         * Provide the tableId.
+         * Provide the table ID.
          *
-         * @param tableId the id for the table.
-         *
-         * @return the builder object.
+         * @param  tableId the id of the Sleeper table
+         * @return         the builder
          */
         public Builder tableId(String tableId) {
             this.tableId = tableId;
@@ -209,9 +229,8 @@ public Builder tableId(String tableId) {
         /**
          * Provide the exportId.
          *
-         * @param exportId the id for the export.
-         *
-         * @return the builder object.
+         * @param  exportId the id for the export
+         * @return          the builder
          */
         public Builder exportId(String exportId) {
             this.exportId = exportId;
@@ -221,9 +240,8 @@ public Builder exportId(String exportId) {
         /**
          * Provide the subExportId.
          *
-         * @param subExportId the id for the sub export.
-         *
-         * @return the builder object.
+         * @param  subExportId the id for the sub export
+         * @return             the builder
          */
         public Builder subExportId(String subExportId) {
             this.subExportId = subExportId;
@@ -233,33 +251,30 @@ public Builder subExportId(String subExportId) {
         /**
          * Provide the regions.
          *
-         * @param regions a list of regions.
-         *
-         * @return the builder object.
+         * @param  regions a list of regions
+         * @return         the builder
          */
         public Builder regions(List<Region> regions) {
             this.regions = regions;
             return this;
         }
 
         /**
-         * Provide the leafPartitionId.
+         * Provide the leaf partition ID.
          *
-         * @param leafPartitionId the id for the leaf partition.
-         *
-         * @return the builder object.
+         * @param  leafPartitionId the id for the leaf partition
+         * @return                 the builder
          */
         public Builder leafPartitionId(String leafPartitionId) {
             this.leafPartitionId = leafPartitionId;
             return this;
         }
 
         /**
-         * Provide the partition regions.
-         *
-         * @param partitionRegion a partition regions.
+         * Provide the partition region.
          *
-         * @return the builder object.
+         * @param  partitionRegion the partition region
+         * @return                 the builder
          */
         public Builder partitionRegion(Region partitionRegion) {
             this.partitionRegion = partitionRegion;
@@ -269,20 +284,14 @@ public Builder partitionRegion(Region partitionRegion) {
         /**
          * Provide the leaf partition files.
          *
-         * @param files the files to be exported.
-         *
-         * @return the builder object.
+         * @param  files the files to be exported
+         * @return       the builder object
          */
         public Builder files(List<String> files) {
             this.files = files;
             return this;
         }
 
-        /**
-         * Builds the BulkExportLeafPartitionQuery.
-         *
-         * @return a BulkExportLeafPartitionQuery object.
-         */
         public BulkExportLeafPartitionQuery build() {
             return new BulkExportLeafPartitionQuery(this);
         }

java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/model/BulkExportLeafPartitionQueryJson.java → java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/core/model/BulkExportLeafPartitionQueryJson.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
 import com.google.gson.JsonElement;
 

java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/model/BulkExportLeafPartitionQuerySerDe.java → java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/core/model/BulkExportLeafPartitionQuerySerDe.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;

java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/model/BulkExportQuery.java → java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/core/model/BulkExportQuery.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
 import java.util.Objects;
 import java.util.UUID;

java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/model/BulkExportQuerySerDe.java → java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/core/model/BulkExportQuerySerDe.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;

java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/model/BulkExportQueryValidationException.java → java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/core/model/BulkExportQueryValidationException.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
 import java.util.Optional;
 

java/bulk-export/bulk-export-core/src/main/java/sleeper/bulkexport/core/recordretrieval/BulkExportQuerySplitter.java

@@ -0,0 +1,155 @@
+/*
+ * Copyright 2022-2024 Crown Copyright
+ *
+ * Licensed 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 sleeper.bulkexport.core.recordretrieval;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import sleeper.bulkexport.core.model.BulkExportLeafPartitionQuery;
+import sleeper.bulkexport.core.model.BulkExportQuery;
+import sleeper.core.partition.Partition;
+import sleeper.core.partition.PartitionTree;
+import sleeper.core.properties.table.TableProperties;
+import sleeper.core.statestore.StateStore;
+import sleeper.core.statestore.StateStoreException;
+
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import static sleeper.core.properties.table.TableProperty.QUERY_PROCESSOR_CACHE_TIMEOUT;
+
+/**
+ * Splits up an export query into leaf partition export queries.
+ */
+public class BulkExportQuerySplitter {
+    private static final Logger LOGGER = LoggerFactory.getLogger(BulkExportQuerySplitter.class);
+
+    private final StateStore stateStore;
+    private final TableProperties tableProperties;
+    private final Supplier<String> idSupplier;
+    private final Supplier<Instant> timeSupplier;
+    private List<Partition> leafPartitions;
+    private PartitionTree partitionTree;
+    private Map<String, List<String>> partitionToFiles;
+    private Instant nextInitialiseTime;
+
+    public BulkExportQuerySplitter(TableProperties tableProperties, StateStore stateStore) {
+        this(stateStore, tableProperties, () -> UUID.randomUUID().toString(), Instant::now);
+    }
+
+    public BulkExportQuerySplitter(StateStore stateStore, TableProperties tableProperties, Supplier<String> idSupplier, Supplier<Instant> timeSupplier) {
+        this.stateStore = stateStore;
+        this.tableProperties = tableProperties;
+        this.idSupplier = idSupplier;
+        this.timeSupplier = timeSupplier;
+        init(timeSupplier.get());
+    }
+
+    /**
+     * Initialises the partitions and the mapping from partitions to active files if
+     * needed.
+     *
+     * @throws StateStoreException if the state store can't be accessed
+     */
+    public void initIfNeeded() throws StateStoreException {
+        Instant now = timeSupplier.get();
+        if (nextInitialiseTime.isAfter(now)) {
+            LOGGER.debug("Not refreshing state for table {}", tableProperties.getStatus());
+            return;
+        }
+        init(now);
+    }
+
+    private void init(Instant now) throws StateStoreException {
+        List<Partition> partitions = stateStore.getAllPartitions();
+        Map<String, List<String>> partitionToFileMapping = stateStore.getPartitionToReferencedFilesMap();
+        init(partitions, partitionToFileMapping, now);
+    }
+
+    private void init(List<Partition> partitions, Map<String, List<String>> partitionToFileMapping, Instant now) {
+        leafPartitions = partitions.stream()
+                .filter(Partition::isLeafPartition)
+                .collect(Collectors.toList());
+        partitionTree = new PartitionTree(partitions);
+        partitionToFiles = partitionToFileMapping;
+        nextInitialiseTime = now.plus(tableProperties.getInt(QUERY_PROCESSOR_CACHE_TIMEOUT), ChronoUnit.MINUTES);
+        LOGGER.info("Loaded state for table {}. Found {} partitions. Next initialise time: {}",
+                tableProperties.getStatus(), partitions.size(), nextInitialiseTime);
+    }
+
+    /**
+     * Splits up a query into a sub-query per relevant leaf partition. For each leaf partition, it
+     * finds the parent partitions in the tree and adds any files still belonging
+     * to the parent to the sub query.
+     *
+     * @param  bulkExportQuery the query to be split up
+     * @return                 a list of {@link LeafPartitionQuery}s
+     */
+    public List<BulkExportLeafPartitionQuery> splitIntoLeafPartitionQueries(BulkExportQuery bulkExportQuery) {
+        LOGGER.debug("There are {} relevant leaf partitions", leafPartitions.size());
+
+        List<BulkExportLeafPartitionQuery> leafPartitionQueriesList = new ArrayList<>();
+        for (Partition partition : leafPartitions) {
+            List<String> files = getFiles(partition);
+
+            if (files.isEmpty()) {
+                LOGGER.info("No files for partition {}", partition.getId());
+                continue;
+            }
+
+            // For each leaf partition, create query with pre-populated list of files that
+            // need to be read.
+            BulkExportLeafPartitionQuery bulkExportLeafPartitionQuery = BulkExportLeafPartitionQuery
+                    .forPartition(bulkExportQuery, tableProperties, partition)
+                    .subExportId(idSupplier.get())
+                    .files(files)
+                    .build();
+            LOGGER.debug("Created {}", bulkExportLeafPartitionQuery);
+            leafPartitionQueriesList.add(bulkExportLeafPartitionQuery);
+        }
+
+        return leafPartitionQueriesList;
+    }
+
+    /**
+     * Gets a list of file paths for a partition.
+     *
+     * @param  partition which partition to get the files from
+     * @return           a list of files
+     */
+    protected List<String> getFiles(Partition partition) {
+        // Get all partitions up to the root of the tree
+        List<String> relevantPartitions = new ArrayList<>();
+        relevantPartitions.add(partition.getId());
+        relevantPartitions.addAll(partitionTree.getAllAncestorIds(partition.getId()));
+
+        // Get relevant files
+        List<String> files = new ArrayList<>();
+        for (String partitionId : relevantPartitions) {
+            List<String> filesForPartition = partitionToFiles.get(partitionId);
+            if (null != filesForPartition) {
+                files.addAll(partitionToFiles.get(partitionId));
+            }
+        }
+        return files;
+    }
+}

java/bulk-export/bulk-export-core/src/test/java/sleeper/bulkexport/model/BulkExportLeafPartitionQuerySerDeTest.java → java/bulk-export/bulk-export-core/src/test/java/sleeper/bulkexport/core/model/BulkExportLeafPartitionQuerySerDeTest.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
 import com.google.common.io.CharStreams;
 import org.approvaltests.Approvals;

java/bulk-export/bulk-export-core/src/test/java/sleeper/bulkexport/model/BulkExportLeafPartitionQueryTest.java → java/bulk-export/bulk-export-core/src/test/java/sleeper/bulkexport/core/model/BulkExportLeafPartitionQueryTest.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package sleeper.bulkexport.model;
+package sleeper.bulkexport.core.model;
 
 import org.junit.jupiter.api.Test;