From 9409fea97208fea804f389de34a7bf1e409a2cec Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 22 Sep 2023 12:01:20 +0100 Subject: [PATCH 001/129] Rust compaction update --- .github/config/chunks.yaml | 2 +- .github/workflows/chunk-compaction.yaml | 1 + docs/11-dev-guide.md | 11 + example/full/table.properties | 3 + .../sleeper/compaction/store/StoreUtils.java | 105 ++ .../docker/Dockerfile | 9 +- .../compaction-job-execution/pom.xml | 14 +- .../jobexecution/CompactSortedFiles.java | 255 ++-- .../jobexecution/CompactionMethod.java | 104 ++ java/compaction/compaction-rust/pom.xml | 208 +++ .../compaction/jobexecution/RustBridge.java | 179 +++ .../jobexecution/RustCompaction.java | 119 ++ java/compaction/pom.xml | 1 + .../instance/CompactionProperty.java | 6 + .../properties/table/TableProperty.java | 9 + java/pom.xml | 4 + rust/.cargo/config.toml | 2 + rust/Cargo.toml | 22 + rust/Cross.toml | 6 + rust/compaction/.cargo/config.toml | 2 + rust/compaction/Cargo.toml | 59 + rust/compaction/src/aws_s3.rs | 106 ++ rust/compaction/src/details.rs | 641 +++++++++ rust/compaction/src/lib.rs | 351 +++++ rust/compaction/src/sketch.rs | 392 ++++++ rust/compactor/Cargo.toml | 45 + rust/compactor/src/bin/main.rs | 114 ++ rust/rust_sketch/Cargo.toml | 42 + rust/rust_sketch/build.rs | 50 + .../src/include/byte_array_serializer.hpp | 215 +++ rust/rust_sketch/src/include/common_types.hpp | 30 + rust/rust_sketch/src/include/quantiles.hpp | 479 +++++++ rust/rust_sketch/src/lib.rs | 53 + rust/rust_sketch/src/main.rs | 66 + rust/rust_sketch/src/quantiles.rs | 1249 +++++++++++++++++ scripts/cli/builder/Dockerfile | 13 + scripts/cli/dependencies/Dockerfile | 5 + scripts/cli/runInDocker.sh | 5 + 38 files changed, 4816 insertions(+), 161 deletions(-) create mode 100644 java/compaction/compaction-core/src/main/java/sleeper/compaction/store/StoreUtils.java create mode 100644 java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactionMethod.java create mode 100644 java/compaction/compaction-rust/pom.xml create mode 100644 java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java create mode 100644 java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java create mode 100644 rust/.cargo/config.toml create mode 100644 rust/Cargo.toml create mode 100644 rust/Cross.toml create mode 100644 rust/compaction/.cargo/config.toml create mode 100644 rust/compaction/Cargo.toml create mode 100644 rust/compaction/src/aws_s3.rs create mode 100644 rust/compaction/src/details.rs create mode 100644 rust/compaction/src/lib.rs create mode 100644 rust/compaction/src/sketch.rs create mode 100644 rust/compactor/Cargo.toml create mode 100644 rust/compactor/src/bin/main.rs create mode 100644 rust/rust_sketch/Cargo.toml create mode 100644 rust/rust_sketch/build.rs create mode 100644 rust/rust_sketch/src/include/byte_array_serializer.hpp create mode 100644 rust/rust_sketch/src/include/common_types.hpp create mode 100644 rust/rust_sketch/src/include/quantiles.hpp create mode 100644 rust/rust_sketch/src/lib.rs create mode 100644 rust/rust_sketch/src/main.rs create mode 100644 rust/rust_sketch/src/quantiles.rs diff --git a/.github/config/chunks.yaml b/.github/config/chunks.yaml index 4420b6901c..a869a1342f 100644 --- a/.github/config/chunks.yaml +++ b/.github/config/chunks.yaml @@ -22,7 +22,7 @@ chunks: compaction: name: Compaction workflow: chunk-compaction.yaml - modules: [ compaction, compaction/compaction-job-execution, compaction/compaction-task-creation, compaction/compaction-job-creation, compaction/compaction-status-store, compaction/compaction-core, splitter ] + modules: [ compaction, compaction/compaction-job-execution, compaction/compaction-task-creation, compaction/compaction-job-creation, compaction/compaction-status-store, compaction/compaction-core, compaction/compaction-rust, splitter ] data: name: Data workflow: chunk-data.yaml diff --git a/.github/workflows/chunk-compaction.yaml b/.github/workflows/chunk-compaction.yaml index 4cb805582b..dc37ee6670 100644 --- a/.github/workflows/chunk-compaction.yaml +++ b/.github/workflows/chunk-compaction.yaml @@ -11,6 +11,7 @@ on: - 'java/compaction/pom.xml' - 'java/compaction/compaction-job-execution/**' - 'java/compaction/compaction-task-creation/**' + - 'java/compaction/compaction-rust/**' - 'java/compaction/compaction-job-creation/**' - 'java/compaction/compaction-status-store/**' - 'java/compaction/compaction-core/**' diff --git a/docs/11-dev-guide.md b/docs/11-dev-guide.md index 4946442644..0a91a75268 100644 --- a/docs/11-dev-guide.md +++ b/docs/11-dev-guide.md @@ -20,6 +20,8 @@ You will need the following software: * [Java 11/17](https://openjdk.java.net/install/) * [Maven](https://maven.apache.org/): Tested with v3.8.6 * [NodeJS / NPM](https://github.com/nvm-sh/nvm#installing-and-updating): Tested with NodeJS v16.16.0 and npm v8.11.0 +* [Rust](https://rustup.rs/): Tested with Rust v1.71 +* [Cross-rs](https://github.com/cross-rs/cross) You can use the [Nix package manager](https://nixos.org/download.html) to get up to date versions of all of these. When you have Nix installed, an easy way to get a development environment is to run `nix-shell` at the root of the Sleeper @@ -57,6 +59,15 @@ cd java mvn clean install -Pquick ``` +### Disabling Rust component + +You can disable the building of the Rust modules with: + +```bash +cd java +mvn clean install -Pquick -DskipRust=true +``` + ## Using the codebase The codebase is structured around the components explained in the [design document](12-design.md). The elements of the diff --git a/example/full/table.properties b/example/full/table.properties index db1c3541cc..eaed6a8003 100644 --- a/example/full/table.properties +++ b/example/full/table.properties @@ -98,6 +98,9 @@ sleeper.table.compaction.strategy.sizeratio.ratio=3 # concurrently per partition. sleeper.table.compaction.strategy.sizeratio.max.concurrent.jobs.per.partition=2147483647 +# Select what compation method to use on a table. Current options are JAVA and RUST. Rust compaction support is experimental. +sleeper.table.compaction.method=JAVA + ## The following table properties relate to storing and retrieving metadata for tables. diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/store/StoreUtils.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/store/StoreUtils.java new file mode 100644 index 0000000000..e35619f4ca --- /dev/null +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/store/StoreUtils.java @@ -0,0 +1,105 @@ +/* + * Copyright 2022-2023 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.compaction.store; + +import org.apache.commons.lang3.tuple.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import sleeper.core.key.Key; +import sleeper.core.schema.type.PrimitiveType; +import sleeper.core.statestore.FileInfo; +import sleeper.core.statestore.StateStore; +import sleeper.core.statestore.StateStoreException; + +import java.util.ArrayList; +import java.util.List; + +public class StoreUtils { + + private static final Logger LOGGER = LoggerFactory.getLogger(StoreUtils.class); + + private StoreUtils() { + } + + public static boolean updateStateStoreSuccess(List inputFiles, String outputFile, + String partitionId, long recordsWritten, Object minRowKey0, Object maxRowKey0, + long finishTime, StateStore stateStore, List rowKeyTypes) { + List filesToBeMarkedReadyForGC = new ArrayList<>(); + for (String file : inputFiles) { + FileInfo fileInfo = FileInfo.builder().rowKeyTypes(rowKeyTypes).filename(file) + .partitionId(partitionId).lastStateStoreUpdateTime(finishTime) + .fileStatus(FileInfo.FileStatus.ACTIVE).build(); + filesToBeMarkedReadyForGC.add(fileInfo); + } + FileInfo fileInfo = FileInfo.builder().rowKeyTypes(rowKeyTypes).filename(outputFile) + .partitionId(partitionId).fileStatus(FileInfo.FileStatus.ACTIVE) + .numberOfRecords(recordsWritten) + .minRowKey(recordsWritten > 0 ? Key.create(minRowKey0) : null) + .maxRowKey(recordsWritten > 0 ? Key.create(maxRowKey0) : null) + .lastStateStoreUpdateTime(finishTime).build(); + try { + stateStore.atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFile( + filesToBeMarkedReadyForGC, fileInfo); + LOGGER.debug( + "Called atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFile method on DynamoDBStateStore"); + return true; + } catch (StateStoreException e) { + LOGGER.error( + "Exception updating DynamoDB (moving input files to ready for GC and creating new active file): {}", + e.getMessage()); + return false; + } + } + + public static boolean updateStateStoreSuccess(List inputFiles, + Pair outputFiles, String partition, List childPartitions, + Pair recordsWritten, Pair minKeys, + Pair maxKeys, long finishTime, StateStore stateStore, + List rowKeyTypes) { + List filesToBeMarkedReadyForGC = new ArrayList<>(); + for (String file : inputFiles) { + FileInfo fileInfo = FileInfo.builder().rowKeyTypes(rowKeyTypes).filename(file) + .partitionId(partition).lastStateStoreUpdateTime(finishTime) + .fileStatus(FileInfo.FileStatus.ACTIVE).build(); + filesToBeMarkedReadyForGC.add(fileInfo); + } + FileInfo leftFileInfo = FileInfo.builder().rowKeyTypes(rowKeyTypes) + .filename(outputFiles.getLeft()).partitionId(childPartitions.get(0)) + .fileStatus(FileInfo.FileStatus.ACTIVE).numberOfRecords(recordsWritten.getLeft()) + .minRowKey(recordsWritten.getLeft() > 0 ? Key.create(minKeys.getLeft()) : null) + .maxRowKey(recordsWritten.getLeft() > 0 ? Key.create(maxKeys.getLeft()) : null) + .lastStateStoreUpdateTime(finishTime).build(); + FileInfo rightFileInfo = FileInfo.builder().rowKeyTypes(rowKeyTypes) + .filename(outputFiles.getRight()).partitionId(childPartitions.get(1)) + .fileStatus(FileInfo.FileStatus.ACTIVE).numberOfRecords(recordsWritten.getRight()) + .minRowKey(recordsWritten.getRight() > 0 ? Key.create(minKeys.getRight()) : null) + .maxRowKey(recordsWritten.getRight() > 0 ? Key.create(maxKeys.getRight()) : null) + .lastStateStoreUpdateTime(finishTime).build(); + try { + stateStore.atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFiles( + filesToBeMarkedReadyForGC, leftFileInfo, rightFileInfo); + LOGGER.debug( + "Called atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFile method on DynamoDBStateStore"); + return true; + } catch (StateStoreException e) { + LOGGER.error( + "Exception updating DynamoDB while moving input files to ready for GC and creating new active file", + e); + return false; + } + } +} diff --git a/java/compaction/compaction-job-execution/docker/Dockerfile b/java/compaction/compaction-job-execution/docker/Dockerfile index 34a40e8443..c8df14b776 100644 --- a/java/compaction/compaction-job-execution/docker/Dockerfile +++ b/java/compaction/compaction-job-execution/docker/Dockerfile @@ -11,7 +11,14 @@ # 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. -FROM amazoncorretto:11 +FROM ubuntu:22.04 + +RUN apt-get update +RUN apt-get install -y wget gnupg +RUN wget -O- https://apt.corretto.aws/corretto.key | apt-key add - +RUN echo 'deb https://apt.corretto.aws stable main' | tee /etc/apt/sources.list.d/corretto.list +RUN apt-get update +RUN apt-get install -y java-11-amazon-corretto-jdk COPY compaction-job-execution.jar /compaction-job-execution.jar COPY run.sh /run.sh diff --git a/java/compaction/compaction-job-execution/pom.xml b/java/compaction/compaction-job-execution/pom.xml index 44896d1899..218b070bf7 100644 --- a/java/compaction/compaction-job-execution/pom.xml +++ b/java/compaction/compaction-job-execution/pom.xml @@ -14,8 +14,9 @@ ~ See the License for the specific language governing permissions and ~ limitations under the License. --> - + compaction sleeper @@ -51,6 +52,11 @@ compaction-status-store ${project.parent.version} + + sleeper + compaction-rust + ${project.parent.version} + sleeper common-job @@ -139,11 +145,11 @@ - + org.apache.maven.plugins maven-shade-plugin - + \ No newline at end of file diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactSortedFiles.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactSortedFiles.java index 06f01554bf..2833d651e8 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactSortedFiles.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactSortedFiles.java @@ -15,10 +15,10 @@ */ package sleeper.compaction.jobexecution; +import com.amazonaws.services.sqs.model.UnsupportedOperationException; import com.facebook.collections.ByteArray; import org.apache.commons.io.FilenameUtils; import org.apache.commons.lang3.tuple.ImmutablePair; -import org.apache.commons.lang3.tuple.Pair; import org.apache.datasketches.quantiles.ItemsSketch; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; @@ -29,6 +29,7 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionJobStatusStore; +import sleeper.compaction.store.StoreUtils; import sleeper.configuration.jars.ObjectFactory; import sleeper.configuration.jars.ObjectFactoryException; import sleeper.configuration.properties.instance.InstanceProperties; @@ -37,7 +38,6 @@ import sleeper.core.iterator.IteratorException; import sleeper.core.iterator.MergingIterator; import sleeper.core.iterator.SortedRecordIterator; -import sleeper.core.key.Key; import sleeper.core.record.Record; import sleeper.core.record.SingleKeyComparator; import sleeper.core.record.process.RecordsProcessed; @@ -45,10 +45,7 @@ import sleeper.core.schema.Field; import sleeper.core.schema.Schema; import sleeper.core.schema.type.ByteArrayType; -import sleeper.core.schema.type.PrimitiveType; -import sleeper.core.statestore.FileInfo; import sleeper.core.statestore.StateStore; -import sleeper.core.statestore.StateStoreException; import sleeper.io.parquet.record.ParquetReaderIterator; import sleeper.io.parquet.record.ParquetRecordReader; import sleeper.io.parquet.record.ParquetRecordWriterFactory; @@ -68,10 +65,10 @@ import static sleeper.core.metrics.MetricsLogger.METRICS_LOGGER; /** - * Executes a compaction {@link CompactionJob}, i.e. compacts N input files into a single - * output file (in the case of a compaction job within a partition) or into - * two output files (in the case of a splitting compaction job which reads - * files in one partition and outputs files to the split partition). + * Executes a compaction {@link CompactionJob}, i.e. compacts N input files into a single output + * file (in the case of a compaction job within a partition) or into two output files (in the case + * of a splitting compaction job which reads files in one partition and outputs files to the split + * partition). */ public class CompactSortedFiles { private final InstanceProperties instanceProperties; @@ -87,12 +84,9 @@ public class CompactSortedFiles { private static final Logger LOGGER = LoggerFactory.getLogger(CompactSortedFiles.class); public CompactSortedFiles(InstanceProperties instanceProperties, - TableProperties tableProperties, - ObjectFactory objectFactory, - CompactionJob compactionJob, - StateStore stateStore, - CompactionJobStatusStore jobStatusStore, - String taskId) { + TableProperties tableProperties, ObjectFactory objectFactory, + CompactionJob compactionJob, StateStore stateStore, + CompactionJobStatusStore jobStatusStore, String taskId) { this.instanceProperties = instanceProperties; this.tableProperties = tableProperties; this.schema = this.tableProperties.getSchema(); @@ -110,11 +104,27 @@ public RecordsProcessedSummary compact() throws IOException, IteratorException { LOGGER.info("Compaction job {}: compaction called at {}", id, startTime); jobStatusStore.jobStarted(compactionJob, startTime, taskId); + CompactionMethod method = CompactionMethod.determineCompactionMethod(this.tableProperties, + this.compactionJob); + RecordsProcessed recordsProcessed; - if (!compactionJob.isSplittingJob()) { - recordsProcessed = compactNoSplitting(); + // If using non-default method, try that and fallback to default if it fails, + // otherwise, just use default + if (method != CompactionMethod.DEFAULT) { + try { + switch (method) { + case RUST: + recordsProcessed = RustCompaction.compact(this.schema, this.tableProperties, this.compactionJob, this.stateStore); + break; + default: + throw new UnsupportedOperationException("compaction method unknown"); + } + } catch (IOException e) { + LOGGER.error("Compaction failed with method " + method.toString(), e); + recordsProcessed = useJavaCompaction(); + } } else { - recordsProcessed = compactSplitting(); + recordsProcessed = useJavaCompaction(); } Instant finishTime = Instant.now(); @@ -129,6 +139,14 @@ public RecordsProcessedSummary compact() throws IOException, IteratorException { return summary; } + private RecordsProcessed useJavaCompaction() throws IOException, IteratorException { + if (compactionJob.isSplittingJob()) { + return compactSplitting(); + } else { + return compactNoSplitting(); + } + } + private RecordsProcessed compactNoSplitting() throws IOException, IteratorException { Configuration conf = getConfiguration(); @@ -141,13 +159,16 @@ private RecordsProcessed compactNoSplitting() throws IOException, IteratorExcept // Create writer LOGGER.debug("Creating writer for file {}", compactionJob.getOutputFile()); Path outputPath = new Path(compactionJob.getOutputFile()); - ParquetWriter writer = ParquetRecordWriterFactory.createParquetRecordWriter(outputPath, tableProperties, conf); + ParquetWriter writer = ParquetRecordWriterFactory + .createParquetRecordWriter(outputPath, tableProperties, conf); - LOGGER.info("Compaction job {}: Created writer for file {}", compactionJob.getId(), compactionJob.getOutputFile()); + LOGGER.info("Compaction job {}: Created writer for file {}", compactionJob.getId(), + compactionJob.getOutputFile()); Map keyFieldToSketch = getSketches(); long recordsWritten = 0L; - // Record min and max of the first dimension of the row key (the min is from the first record, the max is from + // Record min and max of the first dimension of the row key (the min is from the first + // record, the max is from // the last). Object minKey = null; Object maxKey = null; @@ -163,7 +184,8 @@ private RecordsProcessed compactNoSplitting() throws IOException, IteratorExcept writer.write(record); recordsWritten++; if (0 == recordsWritten % 1_000_000) { - LOGGER.info("Compaction job {}: Written {} records", compactionJob.getId(), recordsWritten); + LOGGER.info("Compaction job {}: Written {} records", compactionJob.getId(), + recordsWritten); } } writer.close(); @@ -174,8 +196,10 @@ private RecordsProcessed compactNoSplitting() throws IOException, IteratorExcept sketchesFilename = FilenameUtils.removeExtension(sketchesFilename); sketchesFilename = sketchesFilename + ".sketches"; Path sketchesPath = new Path(sketchesFilename); - new SketchesSerDeToS3(schema).saveToHadoopFS(sketchesPath, new Sketches(keyFieldToSketch), conf); - LOGGER.info("Compaction job {}: Wrote sketches file to {}", compactionJob.getId(), sketchesPath); + new SketchesSerDeToS3(schema).saveToHadoopFS(sketchesPath, new Sketches(keyFieldToSketch), + conf); + LOGGER.info("Compaction job {}: Wrote sketches file to {}", compactionJob.getId(), + sketchesPath); for (CloseableIterator iterator : inputIterators) { iterator.close(); @@ -190,16 +214,9 @@ private RecordsProcessed compactNoSplitting() throws IOException, IteratorExcept LOGGER.info("Compaction job {}: Read {} records and wrote {} records", compactionJob.getId(), totalNumberOfRecordsRead, recordsWritten); - updateStateStoreSuccess(compactionJob.getInputFiles(), - compactionJob.getOutputFile(), - compactionJob.getPartitionId(), - recordsWritten, - minKey, - maxKey, - finishTime, - stateStore, - schema.getRowKeyTypes()); - LOGGER.info("Compaction job {}: compaction committed to state store at {}", compactionJob.getId(), LocalDateTime.now()); + StoreUtils.updateStateStoreSuccess(compactionJob.getInputFiles(), compactionJob.getOutputFile(), compactionJob.getPartitionId(), recordsWritten, minKey, maxKey, finishTime, + stateStore, schema.getRowKeyTypes()); + LOGGER.info("Compaction job {}: compaction finished at {}", compactionJob.getId(), LocalDateTime.now()); return new RecordsProcessed(totalNumberOfRecordsRead, recordsWritten); } @@ -215,34 +232,43 @@ private RecordsProcessed compactSplitting() throws IOException, IteratorExceptio // Create writers Path leftPath = new Path(compactionJob.getOutputFiles().getLeft()); - ParquetWriter leftWriter = ParquetRecordWriterFactory.createParquetRecordWriter(leftPath, tableProperties, conf); - LOGGER.debug("Compaction job {}: Created writer for file {}", compactionJob.getId(), compactionJob.getOutputFiles().getLeft()); + ParquetWriter leftWriter = ParquetRecordWriterFactory + .createParquetRecordWriter(leftPath, tableProperties, conf); + LOGGER.debug("Compaction job {}: Created writer for file {}", compactionJob.getId(), + compactionJob.getOutputFiles().getLeft()); Path rightPath = new Path(compactionJob.getOutputFiles().getRight()); - ParquetWriter rightWriter = ParquetRecordWriterFactory.createParquetRecordWriter(rightPath, tableProperties, conf); - LOGGER.debug("Compaction job {}: Created writer for file {}", compactionJob.getId(), compactionJob.getOutputFiles().getRight()); + ParquetWriter rightWriter = ParquetRecordWriterFactory + .createParquetRecordWriter(rightPath, tableProperties, conf); + LOGGER.debug("Compaction job {}: Created writer for file {}", compactionJob.getId(), + compactionJob.getOutputFiles().getRight()); Map leftKeyFieldToSketch = getSketches(); Map rightKeyFieldToSketch = getSketches(); long recordsWrittenToLeftFile = 0L; long recordsWrittenToRightFile = 0L; - // Record min and max of the first dimension of the row key (the min is from the first record, the max is from + // Record min and max of the first dimension of the row key (the min is from the first + // record, the max is from // the last) from both files. Object minKeyLeftFile = null; Object minKeyRightFile = null; Object maxKeyLeftFile = null; Object maxKeyRightFile = null; int dimension = compactionJob.getDimension(); - // Compare using the key of dimension compactionJob.getDimension(), i.e. of that position in the list - SingleKeyComparator keyComparator = new SingleKeyComparator(schema.getRowKeyTypes().get(dimension)); + // Compare using the key of dimension compactionJob.getDimension(), i.e. of that position in + // the list + SingleKeyComparator keyComparator = new SingleKeyComparator( + schema.getRowKeyTypes().get(dimension)); String comparisonKeyFieldName = schema.getRowKeyFieldNames().get(dimension); - LOGGER.debug("Splitting on dimension {} (field name {})", dimension, comparisonKeyFieldName); + LOGGER.debug("Splitting on dimension {} (field name {})", dimension, + comparisonKeyFieldName); Object splitPoint = compactionJob.getSplitPoint(); LOGGER.info("Split point is " + splitPoint); - // TODO This is unnecessarily complicated as the records for the left file will all be written in one go, + // TODO This is unnecessarily complicated as the records for the left file will all be + // written in one go, // followed by the records to the right file. while (mergingIterator.hasNext()) { Record record = mergingIterator.next(); @@ -265,8 +291,10 @@ private RecordsProcessed compactSplitting() throws IOException, IteratorExceptio } if ((recordsWrittenToLeftFile > 0 && 0 == recordsWrittenToLeftFile % 1_000_000) - || (recordsWrittenToRightFile > 0 && 0 == recordsWrittenToRightFile % 1_000_000)) { - LOGGER.info("Compaction job {}: Written {} records to left file and {} records to right file", + || (recordsWrittenToRightFile > 0 + && 0 == recordsWrittenToRightFile % 1_000_000)) { + LOGGER.info( + "Compaction job {}: Written {} records to left file and {} records to right file", compactionJob.getId(), recordsWrittenToLeftFile, recordsWrittenToRightFile); } @@ -280,13 +308,15 @@ private RecordsProcessed compactSplitting() throws IOException, IteratorExceptio leftSketchesFilename = FilenameUtils.removeExtension(leftSketchesFilename); leftSketchesFilename = leftSketchesFilename + ".sketches"; Path leftSketchesPath = new Path(leftSketchesFilename); - new SketchesSerDeToS3(schema).saveToHadoopFS(leftSketchesPath, new Sketches(leftKeyFieldToSketch), conf); + new SketchesSerDeToS3(schema).saveToHadoopFS(leftSketchesPath, + new Sketches(leftKeyFieldToSketch), conf); String rightSketchesFilename = compactionJob.getOutputFiles().getRight(); rightSketchesFilename = FilenameUtils.removeExtension(rightSketchesFilename); rightSketchesFilename = rightSketchesFilename + ".sketches"; Path rightSketchesPath = new Path(rightSketchesFilename); - new SketchesSerDeToS3(schema).saveToHadoopFS(rightSketchesPath, new Sketches(rightKeyFieldToSketch), conf); + new SketchesSerDeToS3(schema).saveToHadoopFS(rightSketchesPath, + new Sketches(rightKeyFieldToSketch), conf); LOGGER.info("Wrote sketches to {} and {}", leftSketchesPath, rightSketchesPath); @@ -302,43 +332,48 @@ private RecordsProcessed compactSplitting() throws IOException, IteratorExceptio } LOGGER.info("Compaction job {}: Read {} records and wrote ({}, {}) records", - compactionJob.getId(), totalNumberOfRecordsRead, recordsWrittenToLeftFile, recordsWrittenToRightFile); + compactionJob.getId(), totalNumberOfRecordsRead, recordsWrittenToLeftFile, + recordsWrittenToRightFile); - updateStateStoreSuccess(compactionJob.getInputFiles(), - compactionJob.getOutputFiles(), - compactionJob.getPartitionId(), - compactionJob.getChildPartitions(), + StoreUtils.updateStateStoreSuccess(compactionJob.getInputFiles(), compactionJob.getOutputFiles(), + compactionJob.getPartitionId(), compactionJob.getChildPartitions(), new ImmutablePair<>(recordsWrittenToLeftFile, recordsWrittenToRightFile), new ImmutablePair<>(minKeyLeftFile, minKeyRightFile), - new ImmutablePair<>(maxKeyLeftFile, maxKeyRightFile), - finishTime, - stateStore, + new ImmutablePair<>(maxKeyLeftFile, maxKeyRightFile), finishTime, stateStore, schema.getRowKeyTypes()); - LOGGER.info("Splitting compaction job {}: compaction committed to state store at {}", compactionJob.getId(), LocalDateTime.now()); - return new RecordsProcessed(totalNumberOfRecordsRead, recordsWrittenToLeftFile + recordsWrittenToRightFile); + LOGGER.info("Compaction job {}: compaction finished at {}", compactionJob.getId(), + LocalDateTime.now()); + return new RecordsProcessed(totalNumberOfRecordsRead, + recordsWrittenToLeftFile + recordsWrittenToRightFile); } - private List> createInputIterators(Configuration conf) throws IOException { + private List> createInputIterators(Configuration conf) + throws IOException { List> inputIterators = new ArrayList<>(); for (String file : compactionJob.getInputFiles()) { - ParquetReader reader = new ParquetRecordReader.Builder(new Path(file), schema).withConf(conf).build(); + ParquetReader reader = new ParquetRecordReader.Builder(new Path(file), schema) + .withConf(conf).build(); ParquetReaderIterator recordIterator = new ParquetReaderIterator(reader); inputIterators.add(recordIterator); - LOGGER.debug("Compaction job {}: Created reader for file {}", compactionJob.getId(), file); + LOGGER.debug("Compaction job {}: Created reader for file {}", compactionJob.getId(), + file); } return inputIterators; } - private CloseableIterator getMergingIterator(List> inputIterators) throws IteratorException { + private CloseableIterator getMergingIterator( + List> inputIterators) throws IteratorException { CloseableIterator mergingIterator = new MergingIterator(schema, inputIterators); // Apply an iterator if one is provided if (null != compactionJob.getIteratorClassName()) { SortedRecordIterator iterator = null; try { - iterator = objectFactory.getObject(compactionJob.getIteratorClassName(), SortedRecordIterator.class); + iterator = objectFactory.getObject(compactionJob.getIteratorClassName(), + SortedRecordIterator.class); } catch (ObjectFactoryException e) { - throw new IteratorException("ObjectFactoryException creating iterator of class " + compactionJob.getIteratorClassName(), e); + throw new IteratorException("ObjectFactoryException creating iterator of class " + + compactionJob.getIteratorClassName(), e); } LOGGER.debug("Created iterator of class {}", compactionJob.getIteratorClassName()); iterator.init(compactionJob.getIteratorConfig(), schema); @@ -352,98 +387,8 @@ private Configuration getConfiguration() { return HadoopConfigurationProvider.getConfigurationForECS(instanceProperties); } - private static boolean updateStateStoreSuccess(List inputFiles, - String outputFile, - String partitionId, - long recordsWritten, - Object minRowKey0, - Object maxRowKey0, - long finishTime, - StateStore stateStore, - List rowKeyTypes) { - List filesToBeMarkedReadyForGC = new ArrayList<>(); - for (String file : inputFiles) { - FileInfo fileInfo = FileInfo.builder() - .rowKeyTypes(rowKeyTypes) - .filename(file) - .partitionId(partitionId) - .lastStateStoreUpdateTime(finishTime) - .fileStatus(FileInfo.FileStatus.ACTIVE) - .build(); - filesToBeMarkedReadyForGC.add(fileInfo); - } - FileInfo fileInfo = FileInfo.builder() - .rowKeyTypes(rowKeyTypes) - .filename(outputFile) - .partitionId(partitionId) - .fileStatus(FileInfo.FileStatus.ACTIVE) - .numberOfRecords(recordsWritten) - .minRowKey(recordsWritten > 0 ? Key.create(minRowKey0) : null) - .maxRowKey(recordsWritten > 0 ? Key.create(maxRowKey0) : null) - .lastStateStoreUpdateTime(finishTime) - .build(); - try { - stateStore.atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFile(filesToBeMarkedReadyForGC, fileInfo); - LOGGER.debug("Called atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFile method on DynamoDBStateStore"); - return true; - } catch (StateStoreException e) { - LOGGER.error("Exception updating DynamoDB (moving input files to ready for GC and creating new active file): {}", e.getMessage()); - return false; - } - } - - private static boolean updateStateStoreSuccess(List inputFiles, - Pair outputFiles, - String partition, - List childPartitions, - Pair recordsWritten, - Pair minKeys, - Pair maxKeys, - long finishTime, - StateStore stateStore, - List rowKeyTypes) { - List filesToBeMarkedReadyForGC = new ArrayList<>(); - for (String file : inputFiles) { - FileInfo fileInfo = FileInfo.builder() - .rowKeyTypes(rowKeyTypes) - .filename(file) - .partitionId(partition) - .lastStateStoreUpdateTime(finishTime) - .fileStatus(FileInfo.FileStatus.ACTIVE) - .build(); - filesToBeMarkedReadyForGC.add(fileInfo); - } - FileInfo leftFileInfo = FileInfo.builder() - .rowKeyTypes(rowKeyTypes) - .filename(outputFiles.getLeft()) - .partitionId(childPartitions.get(0)) - .fileStatus(FileInfo.FileStatus.ACTIVE) - .numberOfRecords(recordsWritten.getLeft()) - .minRowKey(recordsWritten.getLeft() > 0 ? Key.create(minKeys.getLeft()) : null) - .maxRowKey(recordsWritten.getLeft() > 0 ? Key.create(maxKeys.getLeft()) : null) - .lastStateStoreUpdateTime(finishTime) - .build(); - FileInfo rightFileInfo = FileInfo.builder() - .rowKeyTypes(rowKeyTypes) - .filename(outputFiles.getRight()) - .partitionId(childPartitions.get(1)) - .fileStatus(FileInfo.FileStatus.ACTIVE) - .numberOfRecords(recordsWritten.getRight()) - .minRowKey(recordsWritten.getRight() > 0 ? Key.create(minKeys.getRight()) : null) - .maxRowKey(recordsWritten.getRight() > 0 ? Key.create(maxKeys.getRight()) : null) - .lastStateStoreUpdateTime(finishTime) - .build(); - try { - stateStore.atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFiles(filesToBeMarkedReadyForGC, leftFileInfo, rightFileInfo); - LOGGER.debug("Called atomicallyUpdateFilesToReadyForGCAndCreateNewActiveFile method on DynamoDBStateStore"); - return true; - } catch (StateStoreException e) { - LOGGER.error("Exception updating DynamoDB while moving input files to ready for GC and creating new active file", e); - return false; - } - } - - // TODO These methods are copies of the same ones in IngestRecordsFromIterator - move to sketches module + // TODO These methods are copies of the same ones in IngestRecordsFromIterator - move to + // sketches module private Map getSketches() { Map keyFieldToSketch = new HashMap<>(); for (Field rowKeyField : schema.getRowKeyFields()) { diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactionMethod.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactionMethod.java new file mode 100644 index 0000000000..34dc6c0d44 --- /dev/null +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/jobexecution/CompactionMethod.java @@ -0,0 +1,104 @@ +/* + * Copyright 2022-2023 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.compaction.jobexecution; + +import sleeper.compaction.job.CompactionJob; +import sleeper.configuration.properties.table.TableProperties; +import sleeper.configuration.properties.table.TableProperty; + +import java.util.Locale; + +/** + * Different compaction methods for Sleeper which support different capabilities and must be + * selected based on need. + */ +public enum CompactionMethod { + /** Pure Java compaction implementation. */ + JAVA { + @Override + public boolean supportsIterators() { + return true; + } + + @Override + public boolean supportsSplittingCompactions() { + return true; + } + }, + /** + * Rust accelerated compaction method. This uses a native library written in Rust to perform a + * compaction. + */ + RUST { + @Override + public boolean supportsIterators() { + return false; + } + + @Override + public boolean supportsSplittingCompactions() { + return false; + } + }; + + public static final CompactionMethod DEFAULT = CompactionMethod.JAVA; + + /** + * @return true if this compaction method can be used for splitting compactions + */ + public abstract boolean supportsSplittingCompactions(); + + /** + * @return true if this compaction method supports tables with iterators present + */ + public abstract boolean supportsIterators(); + + /** + * Retrieve the compaction method to use based on the job configuration and the table + * properties. If a compaction method doesn't support all the features needed for this job, then + * an alternative compaction method will be chosen. + * + * @param tableProperties configuration for the table being compacted + * @param job the compaction details + * @return appropriate compaction method + */ + public static CompactionMethod determineCompactionMethod(TableProperties tableProperties, + CompactionJob job) { + // First find desired compaction method + String configMethod = tableProperties.get(TableProperty.COMPACTION_METHOD) + .toUpperCase(Locale.ROOT); + + // Convert to enum value and default to Java + CompactionMethod desired; + try { + desired = CompactionMethod.valueOf(configMethod); + } catch (IllegalArgumentException e) { + desired = CompactionMethod.JAVA; + } + + // Is this a splitting compaction? If yes, but method doesn't support it, then bail + if (job.isSplittingJob() && !desired.supportsSplittingCompactions()) { + return CompactionMethod.JAVA; + } + + // Is an iterator specifed, if so can we support this? + if (job.getIteratorClassName() != null && !desired.supportsIterators()) { + return CompactionMethod.JAVA; + } + + return desired; + } +} diff --git a/java/compaction/compaction-rust/pom.xml b/java/compaction/compaction-rust/pom.xml new file mode 100644 index 0000000000..5c711bf645 --- /dev/null +++ b/java/compaction/compaction-rust/pom.xml @@ -0,0 +1,208 @@ + + + + + compaction + sleeper + 0.20.0-SNAPSHOT + + 4.0.0 + compaction-rust + + + + sleeper + compaction-job-creation + ${project.parent.version} + + + sleeper + common-job + ${project.parent.version} + + + sleeper + core + ${project.parent.version} + + + + com.github.jnr + jnr-ffi + ${jnr.ffi.version} + + + + org.scijava + native-lib-loader + ${scijava.native.version} + + + + sleeper + core + ${project.parent.version} + test-jar + test + + + sleeper + configuration + ${project.parent.version} + test-jar + test + + + sleeper + dynamodb-tools + ${project.parent.version} + test-jar + test + + + sleeper + statestore + ${project.parent.version} + test-jar + test + + + org.testcontainers + localstack + test + + + org.testcontainers + testcontainers + test + + + org.testcontainers + junit-jupiter + test + + + org.mockito + mockito-core + test + + + + + + + org.apache.maven.plugins + maven-resources-plugin + ${resources.plugin.version} + + copy-resources + + + ${project.build.outputDirectory}/natives/ + + + ${project.build.directory}/rust/ + + release/*.dylib + release/*.dll + release/*.so + */release/*.dylib + */release/*.dll + */release/*.so + debug/*.dylib + debug/*.dll + debug/*.so + */debug/*.dylib + */debug/*.dll + */debug/*.so + + false + + + UTF-8 + + + + + + + + + skipRust + !true + + + rust + + + + org.codehaus.mojo + exec-maven-plugin + ${exec.plugin.version} + + + + Invoke Rust Cargo build (Linux x86_64) + generate-resources + + exec + + + cross + ${maven.multiModuleProjectDirectory}/../rust + + build + --release + --target-dir + ${project.build.directory}/rust + --target + x86_64-unknown-linux-gnu + + + + + Invoke Rust Cargo build (Linux Aarch64) + generate-resources + + exec + + + cross + ${maven.multiModuleProjectDirectory}/../rust + + build + --release + --target-dir + ${project.build.directory}/rust + --target + aarch64-unknown-linux-gnu + + + -Ctarget-feature=+lse -Ctarget-cpu=neoverse-n1 + + + + + + + + + + \ No newline at end of file diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java new file mode 100644 index 0000000000..362b6aa50d --- /dev/null +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java @@ -0,0 +1,179 @@ +/* + * Copyright 2022-2023 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.compaction.jobexecution; + +import jnr.ffi.LibraryLoader; +import jnr.ffi.Struct; +import jnr.ffi.annotations.In; +import jnr.ffi.annotations.Out; +import jnr.ffi.types.size_t; +import org.scijava.nativelib.JniExtractor; +import org.scijava.nativelib.NativeLoader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; + +public class RustBridge { + /** + * Native library extraction object. This can extract native libraries from the classpath and + * unpack them to a temporary directory. + */ + private static final JniExtractor EXTRACTOR = NativeLoader.getJniExtractor(); + + /** Paths in the JAR file where a native library may have been placed. */ + private static final String[] LIB_PATHS = {"natives/release", + "natives/x86_64-unknown-linux-gnu/release", "natives/aarch64-unknown-linux-gnu/release", + // Rust debug builds will place libraries in different locations + "natives/x86_64-unknown-linux-gnu/debug", "natives/aarch64-unknown-linux-gnu/debug" }; + + private static final Logger LOGGER = LoggerFactory.getLogger(RustBridge.class); + + private static Compaction nativeCompaction = null; + + /** + * Attempt to load the native compaction library. + * + * The native library will be extracted from the classpath and unpacked to a temporary + * directory. The library is then loaded and linked. Multiple locations are checked in the + * classpath, representing different architectures. Thus, if we attempt to load a library for + * the wrong CPU architecture, loading will fail and the next path will be tried. This way, we + * maintain a single JAR file that can work across multiple CPU architectures. + * + * @return the native compaction object + * @throws IOException if an error occurs during loading or linking the native library + */ + public static synchronized Compaction getRustCompactor() throws IOException { + try { + Compaction nativeLib; + + if (nativeCompaction == null) { + nativeLib = extractAndLink(Compaction.class, "compaction"); + nativeCompaction = nativeLib; + } else { + nativeLib = nativeCompaction; + } + + return nativeLib; + + } catch (UnsatisfiedLinkError err) { + throw (IOException) new IOException().initCause(err); + } + } + + /** + * Loads the named library after extracting it from the classpath. + * + * This function extracts the named library from a JAR on the classpath and attempts to load it + * and bind it to the given interface class. The paths in the array {@link LIB_PATHS} are tried + * in order. If a library is found at a path, this method will attempt to load it. If no library + * is found on the classpath or it can't be loaded (e.g. wrong binary format), the next path + * will be tried. + * + * The library named should be given without platform prefixes, e.g. "foo" will be expanded into + * "libfoo.so" or "foo.dll" as appropriate for this platform. + * + * @param clazz the Java interface type for the native library + * @param libName the library name to extract without platform prefixes. + * @return the absolute extracted path, or null if the library couldn't be found + * @throws IOException if an error occured during file extraction + * @throws UnsatisfiedLinkError if the library could not be found or loaded + */ + public static T extractAndLink(Class clazz, String libName) throws IOException { + // Work through each potential path to see if we can load the library + // successfully + for (String path : LIB_PATHS) { + LOGGER.debug("Attempting to load native library from JAR path {}", path); + // Attempt extraction + File extractedLib = EXTRACTOR.extractJni(path, libName); + + // If file located, attempt to load + if (extractedLib != null) { + LOGGER.debug("Extracted file is at {}", extractedLib); + try { + return LibraryLoader.create(clazz).failImmediately() + .load(extractedLib.getAbsolutePath()); + } catch (UnsatisfiedLinkError e) { + // wrong library, try the next path + LOGGER.error("Unable to load native library from " + path, e); + } + } + } + + // No matches + throw new UnsatisfiedLinkError("Couldn't locate or load " + libName); + } + + /** + * The compaction output data that the native code will populate. + */ + @SuppressWarnings(value = { "checkstyle:membername", "checkstyle:parametername" }) + public static class FFICompactionResult extends Struct { + public final Struct.Pointer min_key = new Struct.Pointer(); + public final Struct.size_t min_key_len = new Struct.size_t(); + + public final Struct.Pointer max_key = new Struct.Pointer(); + public final Struct.size_t max_key_len = new Struct.size_t(); + public final Struct.size_t rows_read = new Struct.size_t(); + public final Struct.size_t rows_written = new Struct.size_t(); + + public FFICompactionResult(jnr.ffi.Runtime runtime) { + super(runtime); + } + + public byte[] getMinKey() { + return toBytes(min_key, min_key_len.intValue()); + } + + public byte[] getMaxKey() { + return toBytes(max_key, max_key_len.intValue()); + } + + public static byte[] toBytes(Struct.Pointer p, int size) { + nullCheck(p); + byte[] result = new byte[size]; + p.get().get(0, result, 0, size); + return result; + } + + public static void nullCheck(Struct.Pointer p) { + if (p.get() == null || p.get().address() == 0) { + throw new IllegalStateException( + "Pointer is null, this struct hasn't been populated?"); + } + } + } + + /** + * The interface for the native library we are calling. + */ + public interface Compaction { + FFICompactionResult allocate_result(); + + void free_result(@In FFICompactionResult res); + + @SuppressWarnings(value = "checkstyle:parametername") + int ffi_merge_sorted_files(@In String[] input_file_paths, @size_t long input_file_paths_len, + String output_file_path, @size_t long row_group_size, @size_t long max_page_size, + @In long[] row_key_columns, @size_t long row_key_columns_len, + @In long[] sort_columns, @size_t long sort_columns_len, + @Out FFICompactionResult result); + } + + private RustBridge() { + } +} diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java new file mode 100644 index 0000000000..dc28a0b8e8 --- /dev/null +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java @@ -0,0 +1,119 @@ +/* + * Copyright 2022-2023 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.compaction.jobexecution; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import sleeper.compaction.job.CompactionJob; +import sleeper.compaction.store.StoreUtils; +import sleeper.configuration.properties.table.TableProperties; +import sleeper.configuration.properties.table.TableProperty; +import sleeper.core.record.process.RecordsProcessed; +import sleeper.core.schema.Schema; +import sleeper.core.schema.type.ByteArrayType; +import sleeper.core.schema.type.IntType; +import sleeper.core.schema.type.LongType; +import sleeper.core.schema.type.PrimitiveType; +import sleeper.core.schema.type.StringType; +import sleeper.core.statestore.StateStore; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.time.LocalDateTime; +import java.util.stream.LongStream; + +public class RustCompaction { + /** Maximum number of rows in a Parquet row group. */ + private static final long RUST_MAX_ROW_GROUP_ROWS = 1_000_000; + + private static final Logger LOGGER = LoggerFactory.getLogger(RustCompaction.class); + + private RustCompaction() { + } + + public static RecordsProcessed compact(Schema schema, TableProperties tableProperties, + CompactionJob compactionJob, StateStore stateStore) throws IOException { + // Obtain native library. This throws an exception if native library can't be loaded and + // linked + RustBridge.Compaction nativeLib = RustBridge.getRustCompactor(); + + // Figure out column numbers for row key fields + // Row keys are numbered from zero, sort keys follow that + long[] rowKeys = LongStream.rangeClosed(0, schema.getRowKeyFields().size()).toArray(); + long[] sortKeys = LongStream + .rangeClosed(0, schema.getRowKeyFields().size() + schema.getSortKeyFields().size()) + .toArray(); + + // Get the page size from table properties + long maxPageSize = tableProperties.getLong(TableProperty.PAGE_SIZE); + + // Create object to hold the result (in native memory) + RustBridge.FFICompactionResult compactionData = nativeLib.allocate_result(); + + try { + // Perform compaction + int result = nativeLib.ffi_merge_sorted_files( + compactionJob.getInputFiles().toArray(new String[0]), + compactionJob.getInputFiles().size(), compactionJob.getOutputFile(), + RUST_MAX_ROW_GROUP_ROWS, maxPageSize, rowKeys, rowKeys.length, sortKeys, + sortKeys.length, compactionData); + + long finishTime = System.currentTimeMillis(); + + // Check result + if (result != 0) { + LOGGER.error("Rust compaction failed, return code: {}", result); + throw new IOException("Rust compaction failed with return code " + result); + } + + long totalNumberOfRecordsRead = compactionData.rows_read.get(); + long recordsWritten = compactionData.rows_written.get(); + + Object minKey = convertBytesToType(compactionData.getMinKey(), schema.getRowKeyTypes().get(0)); + Object maxKey = convertBytesToType(compactionData.getMaxKey(), schema.getRowKeyTypes().get(0)); + + LOGGER.info("Compaction job {}: Read {} records and wrote {} records", + compactionJob.getId(), totalNumberOfRecordsRead, recordsWritten); + + StoreUtils.updateStateStoreSuccess(compactionJob.getInputFiles(), compactionJob.getOutputFile(), + compactionJob.getPartitionId(), recordsWritten, minKey, maxKey, finishTime, + stateStore, schema.getRowKeyTypes()); + LOGGER.info("Compaction job {}: compaction finished at {}", compactionJob.getId(), + LocalDateTime.now()); + + return new RecordsProcessed(totalNumberOfRecordsRead, recordsWritten); + } finally { + // Ensure de-allocation + nativeLib.free_result(compactionData); + } + } + + @SuppressWarnings("unchecked") + private static T convertBytesToType(byte[] arr, PrimitiveType t) { + if (t instanceof IntType) { + return (T) Integer.valueOf(Integer.parseInt(new String(arr, StandardCharsets.UTF_8))); + } else if (t instanceof LongType) { + return (T) Long.valueOf(Long.parseLong(new String(arr, StandardCharsets.UTF_8))); + } else if (t instanceof StringType) { + return (T) new String(arr, StandardCharsets.UTF_8); + } else if (t instanceof ByteArrayType) { + return (T) arr; + } else { + throw new IllegalArgumentException("unrecognised primitive type"); + } + } +} diff --git a/java/compaction/pom.xml b/java/compaction/pom.xml index 88dabd4bc8..dd520c194c 100644 --- a/java/compaction/pom.xml +++ b/java/compaction/pom.xml @@ -30,6 +30,7 @@ compaction-core compaction-job-creation compaction-job-execution + compaction-rust compaction-status-store compaction-task-creation diff --git a/java/configuration/src/main/java/sleeper/configuration/properties/instance/CompactionProperty.java b/java/configuration/src/main/java/sleeper/configuration/properties/instance/CompactionProperty.java index f334af521c..954a291045 100644 --- a/java/configuration/src/main/java/sleeper/configuration/properties/instance/CompactionProperty.java +++ b/java/configuration/src/main/java/sleeper/configuration/properties/instance/CompactionProperty.java @@ -177,6 +177,12 @@ public interface CompactionProperty { .defaultValue("" + Integer.MAX_VALUE) .propertyGroup(InstancePropertyGroup.COMPACTION).build(); + UserDefinedInstanceProperty DEFAULT_COMPACTION_METHOD = Index.propertyBuilder("sleeper.default.table.compaction.method") + .description("Select what compation method to use on a table. Current options are JAVA and RUST. Rust compaction support is" + + "experimental.") + .defaultValue("JAVA") + .propertyGroup(InstancePropertyGroup.COMPACTION).build(); + static List getAll() { return Index.INSTANCE.getAll(); } diff --git a/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java b/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java index 5e466c01c1..96db86ba7d 100644 --- a/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java +++ b/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java @@ -28,6 +28,7 @@ import static sleeper.configuration.Utils.describeEnumValuesInLowerCase; import static sleeper.configuration.properties.instance.CompactionProperty.DEFAULT_COMPACTION_FILES_BATCH_SIZE; +import static sleeper.configuration.properties.instance.CompactionProperty.DEFAULT_COMPACTION_METHOD; import static sleeper.configuration.properties.instance.CompactionProperty.DEFAULT_COMPACTION_STRATEGY_CLASS; import static sleeper.configuration.properties.instance.CompactionProperty.DEFAULT_SIZERATIO_COMPACTION_STRATEGY_MAX_CONCURRENT_JOBS_PER_PARTITION; import static sleeper.configuration.properties.instance.CompactionProperty.DEFAULT_SIZERATIO_COMPACTION_STRATEGY_RATIO; @@ -202,6 +203,14 @@ public interface TableProperty extends SleeperProperty { "concurrently per partition.") .propertyGroup(TablePropertyGroup.COMPACTION) .build(); + + TableProperty COMPACTION_METHOD = Index.propertyBuilder("sleeper.table.compaction.method") + .defaultProperty(DEFAULT_COMPACTION_METHOD) + .description("Select what compation method to use on a table. Current options are JAVA and RUST. Rust compaction support is" + + "experimental.") + .propertyGroup(TablePropertyGroup.COMPACTION) + .build(); + TableProperty STATESTORE_CLASSNAME = Index.propertyBuilder("sleeper.table.statestore.classname") .defaultValue("sleeper.statestore.dynamodb.DynamoDBStateStore") .description("The name of the class used for the metadata store. The default is DynamoDBStateStore. " + diff --git a/java/pom.xml b/java/pom.xml index 19703b7398..65f69f019f 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -150,6 +150,8 @@ 1.2.2 2.3.3 1.4 + 2.2.13 + 2.4.0 5.9.1 4.11.0 @@ -182,6 +184,8 @@ 1.2.8 8.3.1 3.12.1 + 3.1.0 + 3.3.1 + + + + org.apache.maven.plugins + maven-shade-plugin + + org.apache.maven.plugins maven-resources-plugin @@ -175,7 +180,7 @@ - + + + \ No newline at end of file diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java index 290c202fea..ef409b25e4 100644 --- a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java @@ -122,7 +122,9 @@ public static T extractAndLink(Class clazz, String libName) throws IOExce } /** - * The compaction input data that will be populated from the Java side. + * The compaction input data that will be populated from the Java side. If you updated + * this struct (field ordering, types, etc.), you MUST update the corresponding Rust definition + * in rust/compaction/src/lib.rs. */ @SuppressWarnings(value = {"checkstyle:membername"}) @SuppressFBWarnings(value = {"URF_UNREAD_PUBLIC_OR_PROTECTED_FIELD"}) diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java index 9a9bff3d45..c80009dcca 100644 --- a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java @@ -40,6 +40,7 @@ import static sleeper.configuration.properties.table.TableProperty.COMPRESSION_CODEC; import static sleeper.configuration.properties.table.TableProperty.DICTIONARY_ENCODING_FOR_ROW_KEY_FIELDS; import static sleeper.configuration.properties.table.TableProperty.DICTIONARY_ENCODING_FOR_SORT_KEY_FIELDS; +import static sleeper.configuration.properties.table.TableProperty.DICTIONARY_ENCODING_FOR_VALUE_FIELDS; import static sleeper.configuration.properties.table.TableProperty.PAGE_SIZE; import static sleeper.configuration.properties.table.TableProperty.PARQUET_WRITER_VERSION; import static sleeper.configuration.properties.table.TableProperty.STATISTICS_TRUNCATE_LENGTH; @@ -113,8 +114,9 @@ public static FFICompactionParams createFFIParams(CompactionJob job, TableProper params.stats_truncate_length.set(tableProperties.getInt(STATISTICS_TRUNCATE_LENGTH)); params.dict_enc_row_keys.set(tableProperties.getBoolean(DICTIONARY_ENCODING_FOR_ROW_KEY_FIELDS)); params.dict_enc_sort_keys.set(tableProperties.getBoolean(DICTIONARY_ENCODING_FOR_SORT_KEY_FIELDS)); - // Sanity check: minimise lifetimeerties.getBoolean(DICTIONARY_ENCODING_FOR_VALUE_FIELDS)); - { // Sanity check: minimise lifetime + params.dict_enc_values.set(tableProperties.getBoolean(DICTIONARY_ENCODING_FOR_VALUE_FIELDS)); + // Sanity check: minimise lifetime + { Object[] regionMins = region.getRanges().stream().map(Range::getMin).toArray(); params.region_mins.populate(regionMins); } diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 68c7b30060..4b7b425b95 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -43,6 +43,14 @@ use parquet::{ use std::{cell::RefCell, path::PathBuf, sync::Arc}; use url::Url; +/// Type safe variant for Sleeper partition boundary +pub enum RangeBound { + Int32 { val: i32 }, + Int64 { val: i64 }, + String { val: &'static str }, + ByteArray { val: &'static [u8] }, +} + /// A simple iterator for a batch of rows (owned). #[derive(Debug)] struct OwnedRowIter { diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 1033eabdbf..a085478a45 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -30,7 +30,7 @@ use aws_config::BehaviorVersion; use aws_credential_types::provider::ProvideCredentials; use chrono::Local; use futures::TryFutureExt; -use libc::{size_t, EFAULT, EINVAL, EIO}; +use libc::{c_void, size_t, EFAULT, EINVAL, EIO}; use log::{error, info, LevelFilter}; use std::io::Write; use std::sync::Once; @@ -40,9 +40,9 @@ use std::{ }; use url::Url; -// Just publicly expose this function pub use details::merge_sorted_files; pub use details::CompactionResult; +pub use details::RangeBound; /// An object guaranteed to only initialise once. Thread safe. static LOG_CFG: Once = Once::new(); @@ -74,15 +74,6 @@ fn maybe_cfg_log() { }); } -/// Create a vector from a C pointer to a type. -/// -/// # Errors -/// If the array length is invalid, then behaviour is undefined. -#[must_use] -fn array_helper(array: *const T, len: usize) -> Vec { - unsafe { slice::from_raw_parts(array, len).to_vec() } -} - /// Obtains AWS credentials from normal places and then calls merge function /// with obtained credentials. /// @@ -119,6 +110,38 @@ async fn credentials_and_merge( .await } +/// Contains all the input data for setting up a compaction. +/// +/// See java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java +/// for details. Field ordering and types MUST match between the two definitions! +#[repr(C)] +pub struct FFICompactionParams { + input_files_len: usize, + input_files: *const *const c_char, + output_file: *const c_char, + row_key_cols_len: usize, + row_key_cols: *const *const c_char, + sort_key_cols_len: usize, + sort_key_cols: *const *const c_char, + max_row_group_size: usize, + max_page_size: usize, + compression: *const c_char, + writer_version: *const c_char, + column_truncate_length: usize, + stats_truncate_length: usize, + dict_enc_row_keys: bool, + dict_enc_sort_keys: bool, + dict_enc_values: bool, + region_mins_len: usize, + region_mins: *const c_void, + region_maxs_len: usize, + region_maxs: *const c_void, + region_mins_inclusive_len: usize, + region_mins_inclusive: *const *const bool, + region_maxs_inclusive_len: usize, + region_maxs_inclusive: *const *const bool, +} + /// Contains all output data from a compaction operation. #[repr(C)] pub struct FFICompactionResult { @@ -149,8 +172,7 @@ pub extern "C" fn allocate_result() -> *const FFICompactionResult { /// Provides the C FFI interface to calling the [`merge_sorted_files`] function. /// -/// This function has the same signature as [`merge_sorted_files`], but with -/// C FFI bindings. This function validates the pointers are valid strings (or +/// This function takes an `FFICompactionParams` struct which contains all the This function validates the pointers are valid strings (or /// at least attempts to), but undefined behaviour will result if bad pointers are passed. /// /// It is also undefined behaviour to specify and incorrect array length for any array. @@ -180,51 +202,26 @@ pub extern "C" fn allocate_result() -> *const FFICompactionResult { #[allow(clippy::not_unsafe_ptr_arg_deref)] #[no_mangle] pub extern "C" fn ffi_merge_sorted_files( - input_file_paths: *const *const c_char, - input_file_paths_len: size_t, - output_file_path: *const c_char, - row_group_size: size_t, - max_page_size: size_t, - row_key_columns: *const size_t, - row_key_columns_len: size_t, - sort_columns: *const size_t, - sort_columns_len: size_t, + input_data: *mut FFICompactionParams, output_data: *mut FFICompactionResult, ) -> c_int { maybe_cfg_log(); - - // Check for nulls - if input_file_paths.is_null() || output_file_path.is_null() || output_data.is_null() { - error!("Either input or output array pointer or output_data struct pointer is null"); + let Some(params) = (unsafe { input_data.as_ref() }) else { + error!("input data pointer is null"); return EFAULT; - } + }; - // First convert the C string array to an array of Rust string slices. - let Ok(input_paths) = (unsafe { - // create a slice from the pointer - slice::from_raw_parts(input_file_paths, input_file_paths_len) - .iter() - // transform pointer to a non-owned string - .map(|s| { - // is pointer valid? - if s.is_null() { - return Err(ArrowError::InvalidArgumentError(String::new())); - } - // convert to string and check it's valid - CStr::from_ptr(*s) - .to_str() - .map_err(|e| ArrowError::ExternalError(Box::new(e))) - }) - // now convert to a vector if all strings OK, else Err - .map(|x| x.map(Url::parse)) - .collect::, _>>() - }) else { - error!("Error converting input paths as valid UTF-8"); + // Unpack input array pointers + let Ok(input_paths) = unpack_string_array(params.input_files, params.input_files_len) else { + error!("Error converting input paths as valid UTF8"); return EINVAL; }; - // Now unwrap the URL parsing errors - let input_paths = match input_paths.into_iter().collect::, _>>() { + let input_paths = match input_paths + .into_iter() + .map(|x| Url::parse(x)) + .collect::, _>>() + { Ok(v) => v, Err(e) => { error!("URL parsing error on input paths {}", e); @@ -234,48 +231,131 @@ pub extern "C" fn ffi_merge_sorted_files( // Get output file URL let Ok(Ok(output_path)) = - (unsafe { CStr::from_ptr(output_file_path).to_str() }).map(Url::parse) + (unsafe { CStr::from_ptr(params.output_file).to_str() }).map(Url::parse) else { - error!("URL parsing error on output path"); + error!("UTF8 or URL parsing error on output path"); return EINVAL; }; - // Convert C pointer to dynamic arrays - let row_fields = array_helper(row_key_columns, row_key_columns_len); - let sort_cols = array_helper(sort_columns, sort_columns_len); + // Unpack row column names + let Ok(row_key_cols) = unpack_string_array(params.row_key_cols, params.row_key_cols_len) else { + error!("Error converting row kew colun names as valid UTF8"); + return EINVAL; + }; - // Start async runtime - let rt = match tokio::runtime::Runtime::new() { - Ok(v) => v, - Err(e) => { - error!("Couldn't create Rust tokio runtime {}", e); - return EIO; - } + // Unpack sort column names + let Ok(sort_key_cols) = unpack_string_array(params.sort_key_cols, params.sort_key_cols_len) + else { + error!("Error converting sort kew colun names as valid UTF8"); + return EINVAL; }; - // Run compaction - let result = rt.block_on(credentials_and_merge( - &input_paths, - &output_path, - row_group_size, - max_page_size, - &row_fields, - &sort_cols, - )); + // Get compression codec + let Ok(compression_codec) = (unsafe { CStr::from_ptr(params.compression).to_str() }) else { + error!("UTF8 decoding error on compression codec"); + return EINVAL; + }; - match result { - Ok(res) => { - if let Some(data) = unsafe { output_data.as_mut() } { - data.rows_read = res.rows_read; - data.rows_written = res.rows_written; - } - 0 - } - Err(e) => { - error!("merging error {}", e); - -1 - } + // Get writer version + let Ok(writer_version) = (unsafe { CStr::from_ptr(params.writer_version).to_str() }) else { + error!("UTF8 decoding error on writer version"); + return EINVAL; + }; + + let region_mins_inclusive = unpack_bool_array( + params.region_mins_inclusive, + params.region_mins_inclusive_len, + ); + let region_maxs_inclusive = unpack_bool_array( + params.region_maxs_inclusive, + params.region_maxs_inclusive_len, + ); + + println!( + "{:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?}", + input_paths, + output_path, + row_key_cols, + sort_key_cols, + compression_codec, + writer_version, + region_mins_inclusive, + region_maxs_inclusive + ); + // // Start async runtime + // let rt = match tokio::runtime::Runtime::new() { + // Ok(v) => v, + // Err(e) => { + // error!("Couldn't create Rust tokio runtime {}", e); + // return EIO; + // } + // }; + + // // Run compaction + // let result = rt.block_on(credentials_and_merge( + // &input_paths, + // &output_path, + // row_group_size, + // max_page_size, + // &row_fields, + // &sort_cols, + // )); + + // match result { + // Ok(res) => { + // if let Some(data) = unsafe { output_data.as_mut() } { + // data.rows_read = res.rows_read; + // data.rows_written = res.rows_written; + // } + // 0 + // } + // Err(e) => { + // error!("merging error {}", e); + // -1 + // } + // } + let Some(data) = (unsafe { output_data.as_mut() }) else { + error!("output data pointer is null"); + return EFAULT; + }; + data.rows_read = 12; + data.rows_written = 34; + 0 +} + +/// Create a vector from a C pointer to an array of strings. +/// +/// # Errors +/// If the array length is invalid, then behaviour is undefined. +fn unpack_string_array( + array_base: *const *const c_char, + len: usize, +) -> Result, ArrowError> { + unsafe { + // create a slice from the pointer + slice::from_raw_parts(array_base, len) } + .iter() + // transform pointer to a non-owned string + .map(|s| { + // is pointer valid? + if s.is_null() { + return Err(ArrowError::InvalidArgumentError(String::new())); + } + // convert to string and check it's valid + unsafe { CStr::from_ptr(*s) } + .to_str() + .map_err(|e| ArrowError::ExternalError(Box::new(e))) + }) + // now convert to a vector if all strings OK, else Err + .collect::, _>>() +} + +fn unpack_bool_array(array_base: *const *const bool, len: usize) -> Vec { + unsafe { slice::from_raw_parts(array_base, len) } + .iter() + .map(|&bptr| unsafe { *bptr }) + .collect() } /// Free the compaction result previously allocated by [`allocate_result()`]. From 79add9e9a525f3d06cbb4340ceb8fea6c4b2cac8 Mon Sep 17 00:00:00 2001 From: m09526 Date: Mon, 22 Apr 2024 16:49:59 +0000 Subject: [PATCH 041/129] Bridge complete --- .../compaction/jobexecution/RustBridge.java | 25 +++- .../jobexecution/RustCompaction.java | 31 +++++ rust/compaction/src/details.rs | 5 +- rust/compaction/src/lib.rs | 124 ++++++++++++++++-- 4 files changed, 162 insertions(+), 23 deletions(-) diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java index ef409b25e4..dac1ee74a3 100644 --- a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java @@ -135,6 +135,8 @@ public static class FFICompactionParams extends Struct { public final Struct.UTF8StringRef output_file = new Struct.UTF8StringRef(); /** Names of Sleeper row key columns from schema. */ public final Array row_key_cols = new Array<>(this); + /** Types for region schema 1 = Int, 2 = Long, 3 = String, 4 = Byte array. */ + public final Array row_key_schema = new Array<>(this); /** Names of Sleeper sort key columns from schema. */ public final Array sort_key_cols = new Array<>(this); /** Maximum size of output Parquet row group in rows. */ @@ -176,6 +178,7 @@ public FFICompactionParams(jnr.ffi.Runtime runtime) { public void validate() { input_files.validate(); row_key_cols.validate(); + row_key_schema.validate(); sort_key_cols.validate(); region_mins.validate(); region_maxs.validate(); @@ -189,6 +192,9 @@ public void validate() { // Check lengths long rowKeys = row_key_cols.len.get(); + if (rowKeys != row_key_schema.len.get()) { + throw new IllegalStateException("row key schema array has length " + row_key_schema.len.get() + " but there are " + rowKeys + " row key columns"); + } if (rowKeys != region_maxs.len.get()) { throw new IllegalStateException("region maxs has length " + region_maxs.len.get() + " but there are " + rowKeys + " row key columns"); } @@ -304,17 +310,22 @@ protected void setValue(E item, int idx, jnr.ffi.Runtime r) { this.items[idx] = r.getMemoryManager().allocateDirect(r.findType(NativeType.SLONGLONG).size()); this.items[idx].putLong(0, e); } else if (item instanceof java.lang.String) { + // Strings are encoded as 4 byte length then value java.lang.String e = (java.lang.String) item; - byte[] bytes = e.getBytes(StandardCharsets.UTF_8); - // Add one for NULL terminator - int stringSize = bytes.length + 1; - // Allocate memory for string and set it + byte[] utf8string = e.getBytes(StandardCharsets.UTF_8); + // Add four for length + int stringSize = utf8string.length + 4; + // Allocate memory for string and write length then string this.items[idx] = r.getMemoryManager().allocateDirect(stringSize); - this.items[idx].putString(0, e, stringSize, StandardCharsets.UTF_8); + this.items[idx].putInt(0, utf8string.length); + this.items[idx].put(4, utf8string, 0, utf8string.length); } else if (item instanceof byte[]) { + // Byte arrays are encoded as 4 byte length then value byte[] e = (byte[]) item; - this.items[idx] = r.getMemoryManager().allocateDirect(e.length); - this.items[idx].put(0, e, 0, e.length); + int byteSize = e.length + 4; + this.items[idx] = r.getMemoryManager().allocateDirect(byteSize); + this.items[idx].putInt(0, e.length); + this.items[idx].put(4, e, 0, e.length); } else if (item instanceof Boolean) { boolean e = (boolean) item; this.items[idx] = r.getMemoryManager().allocateDirect(1); diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java index c80009dcca..ed337a89af 100644 --- a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java @@ -28,11 +28,17 @@ import sleeper.core.range.Region; import sleeper.core.record.process.RecordsProcessed; import sleeper.core.schema.Schema; +import sleeper.core.schema.type.ByteArrayType; +import sleeper.core.schema.type.IntType; +import sleeper.core.schema.type.LongType; +import sleeper.core.schema.type.PrimitiveType; +import sleeper.core.schema.type.StringType; import sleeper.core.statestore.StateStore; import sleeper.statestore.StateStoreProvider; import java.io.IOException; import java.time.LocalDateTime; +import java.util.List; import java.util.NoSuchElementException; import java.util.Objects; @@ -105,6 +111,7 @@ public static FFICompactionParams createFFIParams(CompactionJob job, TableProper params.input_files.populate(job.getInputFiles().toArray(new String[0])); params.output_file.set(job.getOutputFile()); params.row_key_cols.populate(schema.getRowKeyFieldNames().toArray(new String[0])); + params.row_key_schema.populate(getKeyTypes(schema.getRowKeyTypes())); params.sort_key_cols.populate(schema.getSortKeyFieldNames().toArray(new String[0])); params.max_row_group_size.set(RUST_MAX_ROW_GROUP_ROWS); params.max_page_size.set(tableProperties.getInt(PAGE_SIZE)); @@ -136,6 +143,30 @@ public static FFICompactionParams createFFIParams(CompactionJob job, TableProper return params; } + /** + * Convert a list of Sleeper primitive types to a number indicating their type for FFI translation. + * + * @param keyTypes list of primitive types of columns + * @return array of type IDs + * @throws IllegalStateException if unsupported type found + */ + public static Integer[] getKeyTypes(List keyTypes) { + return keyTypes.stream().mapToInt(type -> { + if (type instanceof IntType) { + return 1; + } else if (type instanceof LongType) { + return 2; + } else if (type instanceof StringType) { + return 3; + } else if (type instanceof ByteArrayType) { + return 4; + } else { + throw new IllegalStateException("Unsupported column type found " + type.getClass()); + } + }).boxed() + .toArray(Integer[]::new); + } + /** * Take the compaction parameters and invoke the Rust compactor using the FFI bridge. * diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 4b7b425b95..c529caff5f 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -44,11 +44,12 @@ use std::{cell::RefCell, path::PathBuf, sync::Arc}; use url::Url; /// Type safe variant for Sleeper partition boundary -pub enum RangeBound { +#[derive(Debug)] +pub enum PartitionBound { Int32 { val: i32 }, Int64 { val: i64 }, String { val: &'static str }, - ByteArray { val: &'static [u8] }, + ByteArray { val: &'static [i8] }, } /// A simple iterator for a batch of rows (owned). diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index a085478a45..3291b51930 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -33,6 +33,7 @@ use futures::TryFutureExt; use libc::{c_void, size_t, EFAULT, EINVAL, EIO}; use log::{error, info, LevelFilter}; use std::io::Write; +use std::str::Utf8Error; use std::sync::Once; use std::{ ffi::{c_char, c_int, CStr}, @@ -42,7 +43,7 @@ use url::Url; pub use details::merge_sorted_files; pub use details::CompactionResult; -pub use details::RangeBound; +pub use details::PartitionBound; /// An object guaranteed to only initialise once. Thread safe. static LOG_CFG: Once = Once::new(); @@ -121,6 +122,8 @@ pub struct FFICompactionParams { output_file: *const c_char, row_key_cols_len: usize, row_key_cols: *const *const c_char, + row_key_schema_len: usize, + row_key_schema: *const *const c_int, sort_key_cols_len: usize, sort_key_cols: *const *const c_char, max_row_group_size: usize, @@ -133,9 +136,9 @@ pub struct FFICompactionParams { dict_enc_sort_keys: bool, dict_enc_values: bool, region_mins_len: usize, - region_mins: *const c_void, + region_mins: *const *const c_void, region_maxs_len: usize, - region_maxs: *const c_void, + region_maxs: *const *const c_void, region_mins_inclusive_len: usize, region_mins_inclusive: *const *const bool, region_maxs_inclusive_len: usize, @@ -262,17 +265,24 @@ pub extern "C" fn ffi_merge_sorted_files( return EINVAL; }; - let region_mins_inclusive = unpack_bool_array( + // Deal with compaction region parameters + let region_mins_inclusive = unpack_primitive_array( params.region_mins_inclusive, params.region_mins_inclusive_len, ); - let region_maxs_inclusive = unpack_bool_array( + let region_maxs_inclusive = unpack_primitive_array( params.region_maxs_inclusive, params.region_maxs_inclusive_len, ); + let schema_types = unpack_primitive_array(params.row_key_schema, params.row_key_schema_len); + let region_mins = + unpack_variant_array(params.region_mins, params.region_mins_len, &schema_types); + let region_maxs = + unpack_variant_array(params.region_maxs, params.region_maxs_len, &schema_types); + println!( - "{:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?}", + "{:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?}", input_paths, output_path, row_key_cols, @@ -280,7 +290,9 @@ pub extern "C" fn ffi_merge_sorted_files( compression_codec, writer_version, region_mins_inclusive, - region_maxs_inclusive + region_maxs_inclusive, + region_mins, + region_maxs, ); // // Start async runtime // let rt = match tokio::runtime::Runtime::new() { @@ -336,28 +348,112 @@ fn unpack_string_array( slice::from_raw_parts(array_base, len) } .iter() + .inspect(|p| { + if p.is_null() { + error!("Found NULL pointer in string array"); + } + }) // transform pointer to a non-owned string .map(|s| { - // is pointer valid? - if s.is_null() { - return Err(ArrowError::InvalidArgumentError(String::new())); + //unpack length (signed because it's from Java) + let str_len = unsafe { *(*s as *const i32) }; + if str_len < 0 { + return Err(ArrowError::InvalidArgumentError(format!( + "Illegal string length in FFI array: {}", + str_len + ))); } // convert to string and check it's valid - unsafe { CStr::from_ptr(*s) } - .to_str() - .map_err(|e| ArrowError::ExternalError(Box::new(e))) + std::str::from_utf8(unsafe { + slice::from_raw_parts(s.byte_add(4) as *const u8, str_len as usize) + }) + .map_err(|e| ArrowError::ExternalError(Box::new(e))) }) // now convert to a vector if all strings OK, else Err .collect::, _>>() } -fn unpack_bool_array(array_base: *const *const bool, len: usize) -> Vec { +/// Create a vector of a primitive array type. +/// +/// # Errors +/// If the array length is invalid, then behaviour is undefined. +fn unpack_primitive_array(array_base: *const *const T, len: usize) -> Vec { unsafe { slice::from_raw_parts(array_base, len) } .iter() + .inspect(|p| { + if p.is_null() { + error!("Found NULL pointer in string array"); + } + }) .map(|&bptr| unsafe { *bptr }) .collect() } +/// Create a vector of a variant array type. Each element may be a +/// i32, i64, String or byte array. The schema types define what decoding is attempted. +/// +/// # Errors +/// If the array length is invalid, then behaviour is undefined. +/// If the schema types are incorrect, then behaviour is undefined. +/// +/// # Panics +/// If the length of the `schema_types` array doesn't match the length specified. +/// +/// Also panics if a negative array length is found in decoding byte arrays or strings. +fn unpack_variant_array( + array_base: *const *const c_void, + len: usize, + schema_types: &[i32], +) -> Result, Utf8Error> { + assert_eq!(len, schema_types.len()); + unsafe { slice::from_raw_parts(array_base, len) } + .iter() + .inspect(|p| { + if p.is_null() { + error!("Found NULL pointer in string array"); + } + }) + .zip(schema_types.iter()) + .map(|(&bptr, type_id)| match type_id { + 1 => Ok(PartitionBound::Int32 { + val: unsafe { *(bptr as *const i32) }, + }), + 2 => Ok(PartitionBound::Int64 { + val: unsafe { *(bptr as *const i64) }, + }), + 3 => { + //unpack length (signed because it's from Java) + let str_len = unsafe { *(bptr as *const i32) }; + if str_len < 0 { + error!("Illegal string length in FFI array: {}", str_len); + panic!("Illegal string length in FFI array: {}", str_len); + } + std::str::from_utf8(unsafe { + slice::from_raw_parts(bptr.byte_add(4) as *const u8, str_len as usize) + }) + .map(|v| PartitionBound::String { val: v }) + } + 4 => { + //unpack length (signed because it's from Java) + let byte_len = unsafe { *(bptr as *const i32) }; + if byte_len < 0 { + error!("Illegal byte array length in FFI array: {}", byte_len); + panic!("Illegal byte array length in FFI array: {}", byte_len); + } + Ok(PartitionBound::ByteArray { + val: unsafe { + slice::from_raw_parts(bptr.byte_add(4) as *const i8, byte_len as usize) + }, + }) + } + x @ _ => { + error!("Unexpected type id {}", x); + panic!("Unexpected type id {}", x); + } + }) + .collect() +} + /// Free the compaction result previously allocated by [`allocate_result()`]. /// /// This function must only be called on pointers to objects allocated by Rust. From b600cda8d84e78e3aa4e9af7907a78827216bb74 Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 23 Apr 2024 13:30:50 +0000 Subject: [PATCH 042/129] Rust side of bridge done --- rust/Cargo.lock | 60 +++++---- rust/compaction/Cargo.toml | 1 + rust/compaction/src/aws_s3.rs | 49 ++++--- rust/compaction/src/details.rs | 53 ++++++-- rust/compaction/src/lib.rs | 239 +++++++++++++++------------------ rust/compactor/src/bin/main.rs | 33 +++-- 6 files changed, 231 insertions(+), 204 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index dc10193d10..86180d9fd0 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -118,6 +118,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "anyhow" +version = "1.0.82" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f538837af36e6f6a9be0faa67f9a314f8119e4e4b5867c6ab40ed60360142519" + [[package]] name = "arrayvec" version = "0.7.4" @@ -389,9 +395,9 @@ checksum = "f1fdabc7756949593fe60f30ec81974b613357de856987752631dea1e3394c80" [[package]] name = "aws-config" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2a89e0000cde82447155d64eeb71720b933b4396a6fbbebad3f8b4f88ca7b54" +checksum = "b2a4707646259764ab59fd9a50e9de2e92c637b28b36285d6f6fa030e915fbd9" dependencies = [ "aws-credential-types", "aws-runtime", @@ -455,9 +461,9 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.20.0" +version = "1.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32fcc572fd5c58489ec205ec3e4e5f7d63018898a485cbf922a462af496bc300" +checksum = "3d70fb493f4183f5102d8a8d0cc9b57aec29a762f55c0e7bf527e0f7177bb408" dependencies = [ "aws-credential-types", "aws-runtime", @@ -477,9 +483,9 @@ dependencies = [ [[package]] name = "aws-sdk-ssooidc" -version = "1.20.0" +version = "1.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b6275fa8684a1192754221173b1f7a7c1260d6b0571cc2b8af09468eb0cffe5" +checksum = "de3f37549b3e38b7ea5efd419d4d7add6ea1e55223506eb0b4fef9d25e7cc90d" dependencies = [ "aws-credential-types", "aws-runtime", @@ -499,9 +505,9 @@ dependencies = [ [[package]] name = "aws-sdk-sts" -version = "1.20.0" +version = "1.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30acd58272fd567e4853c5075d838be1626b59057e0249c9be5a1a7eb13bf70f" +checksum = "3b2ff219a5d4b795cd33251c19dbe9c4b401f2b2cbe513e07c76ada644eaf34e" dependencies = [ "aws-credential-types", "aws-runtime", @@ -522,9 +528,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11d6f29688a4be9895c0ba8bef861ad0c0dac5c15e9618b9b7a6c233990fc263" +checksum = "58b56f1cbe6fd4d0c2573df72868f20ab1c125ca9c9dbce17927a463433a2e57" dependencies = [ "aws-credential-types", "aws-smithy-http", @@ -556,9 +562,9 @@ dependencies = [ [[package]] name = "aws-smithy-http" -version = "0.60.7" +version = "0.60.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f10fa66956f01540051b0aa7ad54574640f748f9839e843442d99b970d3aff9" +checksum = "4a7de001a1b9a25601016d8057ea16e31a45fdca3751304c8edf4ad72e706c08" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -595,9 +601,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.3.0" +version = "1.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de34bcfa1fb3c82a80e252a753db34a6658e07f23d3a5b3fc96919518fa7a3f5" +checksum = "44e7945379821074549168917e89e60630647e186a69243248f08c6d168b975a" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -661,9 +667,9 @@ dependencies = [ [[package]] name = "aws-smithy-xml" -version = "0.60.7" +version = "0.60.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "872c68cf019c0e4afc5de7753c4f7288ce4b71663212771bf5e4542eb9346ca9" +checksum = "d123fbc2a4adc3c301652ba8e149bf4bc1d1725affb9784eb20c953ace06bf55" dependencies = [ "xmlparser", ] @@ -792,12 +798,13 @@ dependencies = [ [[package]] name = "cc" -version = "1.0.94" +version = "1.0.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17f6e324229dc011159fcc089755d1e2e216a90d43a7dea6853ca740b84f35e7" +checksum = "d32a725bc159af97c3e629873bb9f88fb8cf8a4867175f76dc987815ea07c83b" dependencies = [ "jobserver", "libc", + "once_cell", ] [[package]] @@ -908,6 +915,7 @@ checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" name = "compaction" version = "0.1.0" dependencies = [ + "anyhow", "arrow", "aws-config", "aws-credential-types", @@ -1628,9 +1636,9 @@ checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" [[package]] name = "jobserver" -version = "0.1.30" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "685a7d121ee3f65ae4fddd72b25a04bb36b6af81bc0828f7d5434c0fe60fa3a2" +checksum = "d2b099aaa34a9751c5bf0878add70444e1ed2dd73f347be99003d4577277de6e" dependencies = [ "libc", ] @@ -2565,9 +2573,9 @@ dependencies = [ [[package]] name = "signal-hook-registry" -version = "1.4.1" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1" dependencies = [ "libc", ] @@ -2709,18 +2717,18 @@ dependencies = [ [[package]] name = "thiserror" -version = "1.0.58" +version = "1.0.59" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03468839009160513471e86a034bb2c5c0e4baae3b43f79ffc55c4a5427b3297" +checksum = "f0126ad08bff79f29fc3ae6a55cc72352056dfff61e3ff8bb7129476d44b23aa" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.58" +version = "1.0.59" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7" +checksum = "d1cd413b5d558b4c5bf3680e324a6fa5014e7b7c067a51e69dbdf47eb7148b66" dependencies = [ "proc-macro2", "quote", diff --git a/rust/compaction/Cargo.toml b/rust/compaction/Cargo.toml index ac332c36f9..30fdab39df 100644 --- a/rust/compaction/Cargo.toml +++ b/rust/compaction/Cargo.toml @@ -57,3 +57,4 @@ num-format = { version = "0.4.4" } # Nicely formatted numbers tokio-test = { version = "0.4.2" } # Doc tests env_logger = { version = "0.11.3" } chrono = { version = "0.4.26" } # Log helper +anyhow = { version = "1.0.82" } diff --git a/rust/compaction/src/aws_s3.rs b/rust/compaction/src/aws_s3.rs index 883f4c68db..03cc3cd195 100644 --- a/rust/compaction/src/aws_s3.rs +++ b/rust/compaction/src/aws_s3.rs @@ -38,7 +38,6 @@ use object_store::{ CredentialProvider, Error, GetOptions, GetRange, GetResult, ListResult, MultipartId, ObjectMeta, ObjectStore, PutOptions, PutResult, Result, }; -use tokio::io::AsyncWrite; use url::Url; /// A tuple struct to bridge AWS credentials obtained from the [`aws_config`] crate @@ -225,12 +224,18 @@ impl ObjectStore for LoggingObjectStore { fn put_multipart<'life0, 'life1, 'async_trait>( &'life0 self, location: &'life1 Path, - ) -> ::core::pin::Pin< + ) -> Pin< Box< - dyn ::core::future::Future< - Output = Result<(MultipartId, Box)>, - > + ::core::marker::Send - + 'async_trait, + (dyn futures::Future< + Output = std::result::Result< + ( + std::string::String, + Box<(dyn tokio::io::AsyncWrite + std::marker::Send + Unpin + 'static)>, + ), + object_store::Error, + >, + > + std::marker::Send + + 'async_trait), >, > where @@ -241,22 +246,6 @@ impl ObjectStore for LoggingObjectStore { self.store.put_multipart(location) } - fn abort_multipart<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - location: &'life1 Path, - multipart_id: &'life2 MultipartId, - ) -> ::core::pin::Pin< - Box> + ::core::marker::Send + 'async_trait>, - > - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { - self.store.abort_multipart(location, multipart_id) - } - fn get_opts<'life0, 'life1, 'async_trait>( &'life0 self, location: &'life1 Path, @@ -387,6 +376,22 @@ impl ObjectStore for LoggingObjectStore { { self.store.put_opts(location, bytes, opts) } + + fn abort_multipart<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + location: &'life1 Path, + multipart_id: &'life2 MultipartId, + ) -> ::core::pin::Pin< + Box> + ::core::marker::Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + self.store.abort_multipart(location, multipart_id) + } } impl CountingObjectStore for LoggingObjectStore { diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index c529caff5f..1a0a18040b 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -40,11 +40,11 @@ use parquet::{ basic::{Compression, ZstdLevel}, file::properties::WriterProperties, }; -use std::{cell::RefCell, path::PathBuf, sync::Arc}; +use std::{cell::RefCell, collections::HashMap, path::PathBuf, sync::Arc}; use url::Url; /// Type safe variant for Sleeper partition boundary -#[derive(Debug)] +#[derive(Debug, Copy, Clone)] pub enum PartitionBound { Int32 { val: i32 }, Int64 { val: i64 }, @@ -52,6 +52,34 @@ pub enum PartitionBound { ByteArray { val: &'static [i8] }, } +/// All the information for a a Sleeper compaction. +#[derive(Debug)] +pub struct CompactionDetails { + pub input_files: Vec, + pub output_file: Url, + pub row_key_cols: Vec, + pub sort_key_cols: Vec, + pub max_row_group_size: usize, + pub max_page_size: usize, + pub compression: String, + pub writer_version: String, + pub column_truncate_length: usize, + pub stats_truncate_length: usize, + pub dict_enc_row_keys: bool, + pub dict_enc_sort_keys: bool, + pub dict_enc_values: bool, + pub region: HashMap, +} + +/// Defines a partition range of a single column. +#[derive(Debug, Copy, Clone)] +pub struct ColRange { + pub lower: PartitionBound, + pub lower_inclusive: bool, + pub upper: PartitionBound, + pub upper_inclusive: bool, +} + /// A simple iterator for a batch of rows (owned). #[derive(Debug)] struct OwnedRowIter { @@ -135,13 +163,16 @@ impl Iterator for OwnedRowIter { pub async fn merge_sorted_files( aws_creds: Option, region: &Region, - input_file_paths: &[Url], - output_file_path: &Url, - row_group_size: usize, - max_page_size: usize, - row_key_fields: impl AsRef<[usize]>, - sort_columns: impl AsRef<[usize]>, + input_data: &CompactionDetails, ) -> Result { + let input_file_paths = input_data.input_files.clone(); + let output_file_path = input_data.output_file.clone(); + let row_group_size = input_data.max_row_group_size; + let max_page_size = input_data.max_page_size; + let row_key_fields = (0..input_data.row_key_cols.len()).collect::>(); + let sort_columns = + (0..input_data.row_key_cols.len() + input_data.sort_key_cols.len()).collect::>(); + // Read the schema from the first file if input_file_paths.is_empty() { Err(ArrowError::InvalidArgumentError( @@ -173,7 +204,6 @@ pub async fn merge_sorted_files( if validate_schemas_same(&store_factory, &input_file_paths, &schema).await { // Sort the row key column numbers let sorted_row_keys: Vec = row_key_fields - .as_ref() .iter() .sorted() .map(usize::to_owned) @@ -184,14 +214,13 @@ pub async fn merge_sorted_files( // so if schema has 5 columns [0, 1, 2, 3, 4] and sort_columns is [1,4], // then the full list of columns should be [1, 4, 0, 2, 3] let complete_sort_columns: Vec<_> = sort_columns - .as_ref() .iter() .copied() //Deref the integers // then chain to an iterator of [0, schema.len) with sort columns fitered out - .chain((0..schema.fields().len()).filter(|i| !sort_columns.as_ref().contains(i))) + .chain((0..schema.fields().len()).filter(|i| !sort_columns.contains(i))) .collect(); - debug!("Sort columns {:?}", sort_columns.as_ref()); + debug!("Sort columns {:?}", sort_columns); info!("Sorted column order {:?}", complete_sort_columns); // Create a vec of all schema columns for row conversion diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 3291b51930..d642c56db1 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -32,6 +32,8 @@ use chrono::Local; use futures::TryFutureExt; use libc::{c_void, size_t, EFAULT, EINVAL, EIO}; use log::{error, info, LevelFilter}; +use std::borrow::Borrow; +use std::collections::HashMap; use std::io::Write; use std::str::Utf8Error; use std::sync::Once; @@ -42,8 +44,7 @@ use std::{ use url::Url; pub use details::merge_sorted_files; -pub use details::CompactionResult; -pub use details::PartitionBound; +pub use details::{ColRange, CompactionDetails, CompactionResult, PartitionBound}; /// An object guaranteed to only initialise once. Thread safe. static LOG_CFG: Once = Once::new(); @@ -79,12 +80,7 @@ fn maybe_cfg_log() { /// with obtained credentials. /// async fn credentials_and_merge( - input_paths: &[Url], - output_path: &Url, - row_group_size: size_t, - max_page_size: size_t, - row_fields: &[usize], - sort_cols: &[usize], + input_data: &CompactionDetails, ) -> Result { let config = aws_config::defaults(BehaviorVersion::latest()).load().await; let region = config.region().ok_or(ArrowError::InvalidArgumentError( @@ -98,17 +94,7 @@ async fn credentials_and_merge( .provide_credentials() .map_err(|e| ArrowError::ExternalError(Box::new(e))) .await?; - merge_sorted_files( - Some(creds), - region, - input_paths, - output_path, - row_group_size, - max_page_size, - row_fields, - sort_cols, - ) - .await + merge_sorted_files(Some(creds), region, input_data).await } /// Contains all the input data for setting up a compaction. @@ -145,6 +131,81 @@ pub struct FFICompactionParams { region_maxs_inclusive: *const *const bool, } +impl TryFrom<&FFICompactionParams> for CompactionDetails { + type Error = anyhow::Error; + + fn try_from(params: &FFICompactionParams) -> Result { + // We do this separately since we need the values for computing the region + let row_key_cols = unpack_string_array(params.row_key_cols, params.row_key_cols_len)? + .into_iter() + .map(String::from) + .collect(); + let region = compute_region(params, &row_key_cols)?; + + Ok(Self { + input_files: unpack_string_array(params.input_files, params.input_files_len)? + .into_iter() + .map(Url::parse) + .collect::, _>>()?, + output_file: unsafe { CStr::from_ptr(params.output_file) } + .to_str() + .map(Url::parse)??, + row_key_cols, + sort_key_cols: unpack_string_array(params.sort_key_cols, params.sort_key_cols_len)? + .into_iter() + .map(String::from) + .collect(), + max_row_group_size: params.max_row_group_size, + max_page_size: params.max_page_size, + compression: unsafe { CStr::from_ptr(params.compression) } + .to_str()? + .to_owned(), + writer_version: unsafe { CStr::from_ptr(params.writer_version) } + .to_str()? + .to_owned(), + column_truncate_length: params.column_truncate_length, + stats_truncate_length: params.stats_truncate_length, + dict_enc_row_keys: params.dict_enc_row_keys, + dict_enc_sort_keys: params.dict_enc_sort_keys, + dict_enc_values: params.dict_enc_values, + region, + }) + } +} + +fn compute_region>( + params: &FFICompactionParams, + row_key_cols: &Vec, +) -> Result, anyhow::Error> { + let region_mins_inclusive = unpack_primitive_array( + params.region_mins_inclusive, + params.region_mins_inclusive_len, + ); + let region_maxs_inclusive = unpack_primitive_array( + params.region_maxs_inclusive, + params.region_maxs_inclusive_len, + ); + let schema_types = unpack_primitive_array(params.row_key_schema, params.row_key_schema_len); + let region_mins = + unpack_variant_array(params.region_mins, params.region_mins_len, &schema_types)?; + let region_maxs = + unpack_variant_array(params.region_maxs, params.region_maxs_len, &schema_types)?; + + let mut map = HashMap::with_capacity(row_key_cols.len()); + for (idx, row_key) in row_key_cols.iter().enumerate() { + map.insert( + String::from(row_key.borrow()), + ColRange { + lower: region_mins[idx], + lower_inclusive: region_mins_inclusive[idx], + upper: region_maxs[idx], + upper_inclusive: region_maxs_inclusive[idx], + }, + ); + } + Ok(map) +} + /// Contains all output data from a compaction operation. #[repr(C)] pub struct FFICompactionResult { @@ -214,125 +275,41 @@ pub extern "C" fn ffi_merge_sorted_files( return EFAULT; }; - // Unpack input array pointers - let Ok(input_paths) = unpack_string_array(params.input_files, params.input_files_len) else { - error!("Error converting input paths as valid UTF8"); - return EINVAL; - }; - // Now unwrap the URL parsing errors - let input_paths = match input_paths - .into_iter() - .map(|x| Url::parse(x)) - .collect::, _>>() - { + // Start async runtime + let rt = match tokio::runtime::Runtime::new() { Ok(v) => v, Err(e) => { - error!("URL parsing error on input paths {}", e); - return EINVAL; + error!("Couldn't create Rust tokio runtime {}", e); + return EIO; } }; - // Get output file URL - let Ok(Ok(output_path)) = - (unsafe { CStr::from_ptr(params.output_file).to_str() }).map(Url::parse) - else { - error!("UTF8 or URL parsing error on output path"); - return EINVAL; - }; - - // Unpack row column names - let Ok(row_key_cols) = unpack_string_array(params.row_key_cols, params.row_key_cols_len) else { - error!("Error converting row kew colun names as valid UTF8"); - return EINVAL; - }; - - // Unpack sort column names - let Ok(sort_key_cols) = unpack_string_array(params.sort_key_cols, params.sort_key_cols_len) - else { - error!("Error converting sort kew colun names as valid UTF8"); - return EINVAL; - }; - - // Get compression codec - let Ok(compression_codec) = (unsafe { CStr::from_ptr(params.compression).to_str() }) else { - error!("UTF8 decoding error on compression codec"); - return EINVAL; - }; - - // Get writer version - let Ok(writer_version) = (unsafe { CStr::from_ptr(params.writer_version).to_str() }) else { - error!("UTF8 decoding error on writer version"); - return EINVAL; + let details = match TryInto::::try_into(params) { + Ok(d) => d, + Err(e) => { + error!("Couldn't convert compaction input data {}", e); + return EINVAL; + } }; - // Deal with compaction region parameters - let region_mins_inclusive = unpack_primitive_array( - params.region_mins_inclusive, - params.region_mins_inclusive_len, - ); - let region_maxs_inclusive = unpack_primitive_array( - params.region_maxs_inclusive, - params.region_maxs_inclusive_len, - ); - - let schema_types = unpack_primitive_array(params.row_key_schema, params.row_key_schema_len); - let region_mins = - unpack_variant_array(params.region_mins, params.region_mins_len, &schema_types); - let region_maxs = - unpack_variant_array(params.region_maxs, params.region_maxs_len, &schema_types); - - println!( - "{:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?} {:?}", - input_paths, - output_path, - row_key_cols, - sort_key_cols, - compression_codec, - writer_version, - region_mins_inclusive, - region_maxs_inclusive, - region_mins, - region_maxs, - ); - // // Start async runtime - // let rt = match tokio::runtime::Runtime::new() { - // Ok(v) => v, - // Err(e) => { - // error!("Couldn't create Rust tokio runtime {}", e); - // return EIO; - // } - // }; - - // // Run compaction - // let result = rt.block_on(credentials_and_merge( - // &input_paths, - // &output_path, - // row_group_size, - // max_page_size, - // &row_fields, - // &sort_cols, - // )); - - // match result { - // Ok(res) => { - // if let Some(data) = unsafe { output_data.as_mut() } { - // data.rows_read = res.rows_read; - // data.rows_written = res.rows_written; - // } - // 0 - // } - // Err(e) => { - // error!("merging error {}", e); - // -1 - // } - // } - let Some(data) = (unsafe { output_data.as_mut() }) else { - error!("output data pointer is null"); - return EFAULT; - }; - data.rows_read = 12; - data.rows_written = 34; - 0 + // Run compaction + let result = rt.block_on(credentials_and_merge(&details)); + match result { + Ok(res) => { + if let Some(data) = unsafe { output_data.as_mut() } { + data.rows_read = res.rows_read; + data.rows_written = res.rows_written; + } else { + error!("output data pointer is null"); + return EFAULT; + } + 0 + } + Err(e) => { + error!("merging error {}", e); + -1 + } + } } /// Create a vector from a C pointer to an array of strings. diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index e9f22576fe..37dc59104b 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -18,9 +18,9 @@ use aws_config::BehaviorVersion; use aws_credential_types::provider::ProvideCredentials; use chrono::Local; use clap::Parser; -use compaction::merge_sorted_files; +use compaction::{merge_sorted_files, CompactionDetails}; use human_panic::setup_panic; -use std::io::Write; +use std::{collections::HashMap, io::Write}; use url::Url; /// Implements a Sleeper compaction algorithm in Rust. @@ -101,16 +101,23 @@ async fn main() -> color_eyre::Result<()> { .provide_credentials() .await?; - merge_sorted_files( - Some(creds), - region, - &input_urls, - &output_url, - args.max_page_size, - args.row_group_size, - args.row_keys, - args.sort_column, - ) - .await?; + let details = CompactionDetails { + input_files: input_urls, + output_file: output_url, + max_page_size: args.max_page_size, + max_row_group_size: args.row_group_size, + column_truncate_length: 1048576, + stats_truncate_length: 1048576, + compression: "zstd".into(), + writer_version: "2.0".into(), + dict_enc_row_keys: true, + dict_enc_sort_keys: true, + dict_enc_values: true, + region: HashMap::default(), + row_key_cols: vec!["key".into()], + sort_key_cols: vec![], + }; + + merge_sorted_files(Some(creds), region, &details).await?; Ok(()) } From 237b9ac9188ad49db3f111a45ff4503ecf968df1 Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 23 Apr 2024 14:12:28 +0000 Subject: [PATCH 043/129] Clippy warnings --- rust/compaction/src/aws_s3.rs | 2 +- rust/compaction/src/details.rs | 4 +-- rust/compaction/src/lib.rs | 50 +++++++++++++++++-------------- rust/compactor/src/bin/main.rs | 8 ++--- rust/rust_sketch/build.rs | 33 ++++++++------------ rust/rust_sketch/src/lib.rs | 6 ++-- rust/rust_sketch/src/main.rs | 15 +++++----- rust/rust_sketch/src/quantiles.rs | 4 +++ 8 files changed, 62 insertions(+), 60 deletions(-) diff --git a/rust/compaction/src/aws_s3.rs b/rust/compaction/src/aws_s3.rs index 03cc3cd195..4d0c514821 100644 --- a/rust/compaction/src/aws_s3.rs +++ b/rust/compaction/src/aws_s3.rs @@ -263,7 +263,7 @@ impl ObjectStore for LoggingObjectStore { Self: 'async_trait, { if let Some(ref get_range) = options.range { - let range = to_range(&get_range); + let range = to_range(get_range); info!( "GET request byte range {} to {} = {} bytes", range.start.to_formatted_string(&Locale::en), diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 1a0a18040b..a0ea3083a9 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -54,7 +54,7 @@ pub enum PartitionBound { /// All the information for a a Sleeper compaction. #[derive(Debug)] -pub struct CompactionDetails { +pub struct CompactionInput { pub input_files: Vec, pub output_file: Url, pub row_key_cols: Vec, @@ -163,7 +163,7 @@ impl Iterator for OwnedRowIter { pub async fn merge_sorted_files( aws_creds: Option, region: &Region, - input_data: &CompactionDetails, + input_data: &CompactionInput, ) -> Result { let input_file_paths = input_data.input_files.clone(); let output_file_path = input_data.output_file.clone(); diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index d642c56db1..b38454345b 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -44,7 +44,7 @@ use std::{ use url::Url; pub use details::merge_sorted_files; -pub use details::{ColRange, CompactionDetails, CompactionResult, PartitionBound}; +pub use details::{ColRange, CompactionInput, CompactionResult, PartitionBound}; /// An object guaranteed to only initialise once. Thread safe. static LOG_CFG: Once = Once::new(); @@ -80,7 +80,7 @@ fn maybe_cfg_log() { /// with obtained credentials. /// async fn credentials_and_merge( - input_data: &CompactionDetails, + input_data: &CompactionInput, ) -> Result { let config = aws_config::defaults(BehaviorVersion::latest()).load().await; let region = config.region().ok_or(ArrowError::InvalidArgumentError( @@ -131,7 +131,7 @@ pub struct FFICompactionParams { region_maxs_inclusive: *const *const bool, } -impl TryFrom<&FFICompactionParams> for CompactionDetails { +impl TryFrom<&FFICompactionParams> for CompactionInput { type Error = anyhow::Error; fn try_from(params: &FFICompactionParams) -> Result { @@ -139,7 +139,7 @@ impl TryFrom<&FFICompactionParams> for CompactionDetails { let row_key_cols = unpack_string_array(params.row_key_cols, params.row_key_cols_len)? .into_iter() .map(String::from) - .collect(); + .collect::>(); let region = compute_region(params, &row_key_cols)?; Ok(Self { @@ -175,7 +175,7 @@ impl TryFrom<&FFICompactionParams> for CompactionDetails { fn compute_region>( params: &FFICompactionParams, - row_key_cols: &Vec, + row_key_cols: &[T], ) -> Result, anyhow::Error> { let region_mins_inclusive = unpack_primitive_array( params.region_mins_inclusive, @@ -284,7 +284,7 @@ pub extern "C" fn ffi_merge_sorted_files( } }; - let details = match TryInto::::try_into(params) { + let details = match TryInto::::try_into(params) { Ok(d) => d, Err(e) => { error!("Couldn't convert compaction input data {}", e); @@ -333,16 +333,18 @@ fn unpack_string_array( // transform pointer to a non-owned string .map(|s| { //unpack length (signed because it's from Java) - let str_len = unsafe { *(*s as *const i32) }; + // This will have been allocated in Java so alignment will be ok + #[allow(clippy::cast_ptr_alignment)] + let str_len = unsafe { *(*s).cast::() }; if str_len < 0 { return Err(ArrowError::InvalidArgumentError(format!( - "Illegal string length in FFI array: {}", - str_len + "Illegal string length in FFI array: {str_len}" ))); } // convert to string and check it's valid std::str::from_utf8(unsafe { - slice::from_raw_parts(s.byte_add(4) as *const u8, str_len as usize) + #[allow(clippy::cast_sign_loss)] + slice::from_raw_parts(s.byte_add(4).cast::(), str_len as usize) }) .map_err(|e| ArrowError::ExternalError(Box::new(e))) }) @@ -393,39 +395,41 @@ fn unpack_variant_array( .zip(schema_types.iter()) .map(|(&bptr, type_id)| match type_id { 1 => Ok(PartitionBound::Int32 { - val: unsafe { *(bptr as *const i32) }, + val: unsafe { *bptr.cast::() }, }), 2 => Ok(PartitionBound::Int64 { - val: unsafe { *(bptr as *const i64) }, + val: unsafe { *bptr.cast::() }, }), 3 => { //unpack length (signed because it's from Java) - let str_len = unsafe { *(bptr as *const i32) }; + let str_len = unsafe { *bptr.cast::() }; if str_len < 0 { - error!("Illegal string length in FFI array: {}", str_len); - panic!("Illegal string length in FFI array: {}", str_len); + error!("Illegal string length in FFI array: {str_len}"); + panic!("Illegal string length in FFI array: {str_len}"); } std::str::from_utf8(unsafe { - slice::from_raw_parts(bptr.byte_add(4) as *const u8, str_len as usize) + #[allow(clippy::cast_sign_loss)] + slice::from_raw_parts(bptr.byte_add(4).cast::(), str_len as usize) }) .map(|v| PartitionBound::String { val: v }) } 4 => { //unpack length (signed because it's from Java) - let byte_len = unsafe { *(bptr as *const i32) }; + let byte_len = unsafe { *bptr.cast::() }; if byte_len < 0 { - error!("Illegal byte array length in FFI array: {}", byte_len); - panic!("Illegal byte array length in FFI array: {}", byte_len); + error!("Illegal byte array length in FFI array: {byte_len}"); + panic!("Illegal byte array length in FFI array: {byte_len}"); } Ok(PartitionBound::ByteArray { val: unsafe { - slice::from_raw_parts(bptr.byte_add(4) as *const i8, byte_len as usize) + #[allow(clippy::cast_sign_loss)] + slice::from_raw_parts(bptr.byte_add(4).cast::(), byte_len as usize) }, }) } - x @ _ => { - error!("Unexpected type id {}", x); - panic!("Unexpected type id {}", x); + x => { + error!("Unexpected type id {x}"); + panic!("Unexpected type id {x}"); } }) .collect() diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index 37dc59104b..c098b61538 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -18,7 +18,7 @@ use aws_config::BehaviorVersion; use aws_credential_types::provider::ProvideCredentials; use chrono::Local; use clap::Parser; -use compaction::{merge_sorted_files, CompactionDetails}; +use compaction::{merge_sorted_files, CompactionInput}; use human_panic::setup_panic; use std::{collections::HashMap, io::Write}; use url::Url; @@ -101,13 +101,13 @@ async fn main() -> color_eyre::Result<()> { .provide_credentials() .await?; - let details = CompactionDetails { + let details = CompactionInput { input_files: input_urls, output_file: output_url, max_page_size: args.max_page_size, max_row_group_size: args.row_group_size, - column_truncate_length: 1048576, - stats_truncate_length: 1048576, + column_truncate_length: 1_048_576, + stats_truncate_length: 1_048_576, compression: "zstd".into(), writer_version: "2.0".into(), dict_enc_row_keys: true, diff --git a/rust/rust_sketch/build.rs b/rust/rust_sketch/build.rs index 37abf7fc45..aa3fc5cdda 100644 --- a/rust/rust_sketch/build.rs +++ b/rust/rust_sketch/build.rs @@ -30,27 +30,20 @@ fn main() { let tag = std::env::var("RUST_SKETCH_DATASKETCH_TAG").unwrap_or(String::from("5.0.2")); // try to open repository in case it already exists - let _ = match git2::Repository::open(path.join("datasketches-cpp")) { - Ok(repo) => repo, - Err(..) => { - // otherwise clone its repository - println!( - "cargo:warning=Git cloned datasketches-cpp from {} tag {}", - url, tag - ); - let repo = match git2::Repository::clone(&url, path.join("datasketches-cpp")) { - Ok(repo) => repo, - Err(e) => panic!("failed to clone from {} tag {}: {}", url, tag, e), - }; - { - let reference = repo.find_reference(&format!("refs/tags/{}", tag)).unwrap(); - let ob = reference.peel_to_tag().unwrap().into_object(); - repo.checkout_tree(&ob, None).unwrap(); - repo.set_head_detached(ob.id()).unwrap(); - } - repo + if git2::Repository::open(path.join("datasketches-cpp")).is_err() { + // otherwise clone its repository + println!("cargo:warning=Git cloned datasketches-cpp from {url} tag {tag}"); + let repo = match git2::Repository::clone(&url, path.join("datasketches-cpp")) { + Ok(repo) => repo, + Err(e) => panic!("failed to clone from {url} tag {tag}: {e}"), + }; + { + let reference = repo.find_reference(&format!("refs/tags/{tag}")).unwrap(); + let ob = reference.peel_to_tag().unwrap().into_object(); + repo.checkout_tree(&ob, None).unwrap(); + repo.set_head_detached(ob.id()).unwrap(); } - }; + } cxx_build::bridges(vec!["src/quantiles.rs"]) .warnings_into_errors(true) diff --git a/rust/rust_sketch/src/lib.rs b/rust/rust_sketch/src/lib.rs index 7c056a267b..8b22b4c009 100644 --- a/rust/rust_sketch/src/lib.rs +++ b/rust/rust_sketch/src/lib.rs @@ -1,11 +1,11 @@ //! The `rust_sketch` crate provides a Rust interface to some of the functionality of the Apache -//! DataSketches library. We use the datasketches-cpp implementation and provide wrappers for it. +//! `DataSketches` library. We use the datasketches-cpp implementation and provide wrappers for it. //! //! Currently we only have the "Quantiles Sketch" wrapper implemented, but others could be added in //! a similar fashion. //! //! ## Building -//! As part of the build process, this crate needs the Apache DataSketches C++ code which it will +//! As part of the build process, this crate needs the Apache `DataSketches` C++ code which it will //! attempt to Git clone from [https://github.com/apache/datasketches-cpp](https://github.com/apache/datasketches-cpp) by default. If you //! wish to override this location, please set the environment variable `RUST_SKETCH_DATASKETCH_URL` //! to point to a new URL on build: @@ -18,7 +18,7 @@ //! **Note** If this crate has already cloned the datasketches-cpp repo then changing the URL to point //! at a different repository will not trigger a new clone operation, you will need to `cargo clean` the //! build directory first. -//! +//! //! ## Performance //! Based on the assumption that sketches are updated often and read infrequently, the API design //! has been created to allow for quick updates, minimizing copies and trying to do moves instead, diff --git a/rust/rust_sketch/src/main.rs b/rust/rust_sketch/src/main.rs index 749d1f67df..8251265593 100644 --- a/rust/rust_sketch/src/main.rs +++ b/rust/rust_sketch/src/main.rs @@ -32,7 +32,8 @@ fn main() { // Serialise it and then deserialise it let serial = sketch.serialize(0).expect("Serialise fail"); - let recreate = rust_sketch::quantiles::byte::byte_deserialize(&serial).expect("Deserialise fail"); + let recreate = + rust_sketch::quantiles::byte::byte_deserialize(&serial).expect("Deserialise fail"); println!( "Sketch contains {} values, median {} 99th q {}", @@ -44,9 +45,9 @@ fn main() { // Check things match! assert_eq!(sketch.get_k(), recreate.get_k()); assert_eq!(sketch.get_n(), recreate.get_n()); - assert_eq!( - sketch.get_normalized_rank_error(true), - recreate.get_normalized_rank_error(true) + assert!( + (sketch.get_normalized_rank_error(true) - recreate.get_normalized_rank_error(true)).abs() + < f64::EPSILON ); assert_eq!( sketch.get_quantile(0.4, true).unwrap(), @@ -56,8 +57,8 @@ fn main() { sketch.get_min_item().unwrap(), recreate.get_min_item().unwrap() ); - assert_eq!( - sketch.get_rank(&[b'f'], true).unwrap(), - recreate.get_rank(&[b'f'], true).unwrap() + assert!( + (sketch.get_rank(&[b'f'], true).unwrap() - recreate.get_rank(&[b'f'], true).unwrap()).abs() + < f64::EPSILON ); } diff --git a/rust/rust_sketch/src/quantiles.rs b/rust/rust_sketch/src/quantiles.rs index 5f866025b3..003d9f3f4a 100644 --- a/rust/rust_sketch/src/quantiles.rs +++ b/rust/rust_sketch/src/quantiles.rs @@ -70,6 +70,7 @@ */ /// Provides a quantiles sketch for the `i32` primitive type. #[cxx::bridge(namespace = "rust_sketch")] +#[allow(clippy::module_name_repetitions, clippy::must_use_candidate)] pub mod i32 { unsafe extern "C++" { @@ -427,6 +428,7 @@ mod i32_tests { /// Provides a quantiles sketch for the `i64` primitive type. #[cxx::bridge(namespace = "rust_sketch")] +#[allow(clippy::module_name_repetitions, clippy::must_use_candidate)] pub mod i64 { unsafe extern "C++" { include!("rust_sketch/src/include/quantiles.hpp"); @@ -774,6 +776,7 @@ mod i64_tests { /// Provides a quantiles sketch for the `str` primitive type. #[cxx::bridge(namespace = "rust_sketch")] +#[allow(clippy::module_name_repetitions, clippy::must_use_candidate)] pub mod str { unsafe extern "C++" { include!("rust_sketch/src/include/quantiles.hpp"); @@ -1127,6 +1130,7 @@ mod str_tests { /// Provides a quantiles sketch for the [`Vec`] byte array type. #[cxx::bridge(namespace = "rust_sketch")] +#[allow(clippy::module_name_repetitions, clippy::must_use_candidate)] pub mod byte { unsafe extern "C++" { include!("rust_sketch/src/include/quantiles.hpp"); From 69aaf1118f551c1bfdad76571fb2b7bfc36fc65c Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 23 Apr 2024 15:03:50 +0000 Subject: [PATCH 044/129] Remove old code --- rust/compaction/src/details.rs | 470 +-------------------------------- 1 file changed, 8 insertions(+), 462 deletions(-) diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index a0ea3083a9..1971e540bc 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -1,7 +1,7 @@ //! The `internal` module contains the internal functionality and error conditions //! to actually implement the compaction library. /* - * Copyright 2022-2023 Crown Copyright + * 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. @@ -80,25 +80,6 @@ pub struct ColRange { pub upper_inclusive: bool, } -/// A simple iterator for a batch of rows (owned). -#[derive(Debug)] -struct OwnedRowIter { - rows: Result, - error_reported: bool, - index: usize, -} - -impl OwnedRowIter { - /// Create a new iterator over some rows. - fn new(rows: Result) -> Self { - Self { - rows, - error_reported: false, - index: 0, - } - } -} - /// Contains compaction results. /// /// This provides the details of compaction results that Sleeper @@ -111,32 +92,6 @@ pub struct CompactionResult { pub rows_written: usize, } -impl Iterator for OwnedRowIter { - type Item = Result; - - fn next(&mut self) -> Option { - match &self.rows { - Ok(rows_val) => { - if self.index == rows_val.num_rows() { - return None; - } - let row = rows_val.row(self.index).owned(); - self.index += 1; - Some(Ok(row)) - } - Err(e) => { - // This is not really the right way to handle errors inside an iterator - if self.error_reported { - None - } else { - self.error_reported = true; - Some(Err(ArrowError::InvalidArgumentError(e.to_string()))) - } - } - } - } -} - /// Merges the given Parquet files and reads the schema from the first. /// /// This function reads the schema from the first file, then calls @@ -165,16 +120,8 @@ pub async fn merge_sorted_files( region: &Region, input_data: &CompactionInput, ) -> Result { - let input_file_paths = input_data.input_files.clone(); - let output_file_path = input_data.output_file.clone(); - let row_group_size = input_data.max_row_group_size; - let max_page_size = input_data.max_page_size; - let row_key_fields = (0..input_data.row_key_cols.len()).collect::>(); - let sort_columns = - (0..input_data.row_key_cols.len() + input_data.sort_key_cols.len()).collect::>(); - // Read the schema from the first file - if input_file_paths.is_empty() { + if input_data.input_files.is_empty() { Err(ArrowError::InvalidArgumentError( "No input paths supplied".into(), )) @@ -182,7 +129,8 @@ pub async fn merge_sorted_files( // Create our object store factory let store_factory = ObjectStoreFactory::new(aws_creds, region); // Java tends to use s3a:// URI scheme instead of s3:// so map it here - let input_file_paths: Vec = input_file_paths + let input_file_paths: Vec = input_data + .input_files .iter() .map(|u| { let mut t = u.clone(); @@ -194,263 +142,15 @@ pub async fn merge_sorted_files( .collect(); // Change output file scheme - let mut output_file_path = output_file_path.clone(); + let mut output_file_path = input_data.output_file.clone(); if output_file_path.scheme() == "s3a" { let _ = output_file_path.set_scheme("s3"); } - // Read Schema from first file - let schema: Arc = read_schema(&store_factory, &input_file_paths[0]).await?; - // validate all files have same schema - if validate_schemas_same(&store_factory, &input_file_paths, &schema).await { - // Sort the row key column numbers - let sorted_row_keys: Vec = row_key_fields - .iter() - .sorted() - .map(usize::to_owned) - .collect(); - - // Create a list of column numbers. Sort columns should be first, - // followed by others in order - // so if schema has 5 columns [0, 1, 2, 3, 4] and sort_columns is [1,4], - // then the full list of columns should be [1, 4, 0, 2, 3] - let complete_sort_columns: Vec<_> = sort_columns - .iter() - .copied() //Deref the integers - // then chain to an iterator of [0, schema.len) with sort columns fitered out - .chain((0..schema.fields().len()).filter(|i| !sort_columns.contains(i))) - .collect(); - - debug!("Sort columns {:?}", sort_columns); - info!("Sorted column order {:?}", complete_sort_columns); - - // Create a vec of all schema columns for row conversion - let fields = schema - .fields() - .iter() - .map(|field| SortField::new(field.data_type().clone())) - .collect::>(); - - // now re-order this list according to the indexes in complete_sort_columns - let fields: Vec<_> = complete_sort_columns - .iter() - .map(|&i| fields[i].clone()) - .collect(); - - let converter = RowConverter::new(fields)?; - - merge_sorted_files_with_schema( - &store_factory, - &input_file_paths, - &output_file_path, - row_group_size, - max_page_size, - sorted_row_keys, - complete_sort_columns, - &schema, - &Arc::new(RefCell::new(converter)), - ) - .await - } else { - Err(ArrowError::SchemaError( - "Schemas do not match across all input files".into(), - )) - } - } -} - -/// Merges sorted Parquet files into a single output file. They must also have the same schema. -/// -/// The files MUST be sorted for this function to work. This is not checked! -/// -/// The files are read and then k-way merged into a single output Parquet file. -/// -/// # Examples -/// ```ignore -/// let result = merge_sorted_files_with_schema( &vec!["file:///path/to/file1.parquet".into(), "file:///path/to/file2.parquet".into()], "file:///path/to/output.parquet".into(), 65535, 1_000_000, my_schema )?; -/// ``` -/// -/// # Errors -/// The supplied input file paths must have a valid extension: parquet, parq, or pq. -/// -#[allow(clippy::too_many_arguments)] -async fn merge_sorted_files_with_schema( - store_factory: &ObjectStoreFactory, - file_paths: &[Url], - output_file_path: &Url, - row_group_size: usize, - max_page_size: usize, - row_key_fields: impl AsRef<[usize]>, - sort_columns: impl AsRef<[usize]>, - schema: &Arc, - converter_ptr: &Arc>, -) -> Result { - // Are the row key columns defined and valid? - if row_key_fields.as_ref().is_empty() { - return Err(ArrowError::InvalidArgumentError( - "row_key_fields cannot be empty".into(), - )); - } else if sort_columns.as_ref().is_empty() { - return Err(ArrowError::InvalidArgumentError( - "sort_columns cannot be empty".into(), - )); - } else if row_key_fields - .as_ref() - .iter() - .any(|&index| index >= schema.fields().len()) - { - return Err(ArrowError::InvalidArgumentError( - "row_key_fields contains invalid column number for schema (too high?)".into(), - )); - } else if sort_columns - .as_ref() - .iter() - .any(|&index| index >= schema.fields().len()) - { - return Err(ArrowError::InvalidArgumentError( - "sort_columns contains invalid column number for schema (too high?)".into(), - )); - } - - let file_iters = future::join_all( - file_paths - .iter() - .map(|file_path| get_file_iterator(store_factory, file_path, None)), - ) - .await - .into_iter() - .collect::, _>>(); - - let converter_clone = converter_ptr.clone(); - - let batch_iters = file_iters?.into_iter().map(|file_iter| { - file_iter.flat_map(|batch| match batch { - Ok(b) => { - let rows = converter_ptr.borrow_mut().convert_columns(b.columns()); - OwnedRowIter::new(rows) - } - Err(e) => OwnedRowIter::new(Err(e.into())), + Ok(CompactionResult { + rows_read: 0, + rows_written: 0, }) - }); - - merge_sorted_iters_into_parquet_file( - store_factory, - batch_iters, - output_file_path, - row_group_size, - max_page_size, - row_key_fields, - schema, - &converter_clone, - ) - .await -} - -/// Takes a group of Arrow record batch iterators and writes the merged -/// output to a Parquet file. -#[allow(clippy::too_many_arguments)] -async fn merge_sorted_iters_into_parquet_file( - store_factory: &ObjectStoreFactory, - iters: A, - output_file_path: &Url, - row_group_size: usize, - max_page_size: usize, - row_key_fields: impl AsRef<[usize]>, - schema: &Arc, - converter_ptr: &Arc>, -) -> Result -where - A: Iterator, - B: Iterator>, -{ - // Create Parquet writer options - let props = WriterProperties::builder() - .set_compression(Compression::ZSTD(ZstdLevel::default())) - .set_data_page_size_limit(max_page_size) - .set_max_row_group_size(row_group_size) - .build(); - - info!("Row_key columns {:?}", row_key_fields.as_ref()); - - // Create Object store async writer - let store = store_factory.get_object_store(output_file_path)?; - let path = object_store::path::Path::from(output_file_path.path()); - let (_id, store_writer) = store - .put_multipart(&path) - .await - .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; - - info!("Writing to {}", output_file_path); - - // Create Parquet writer - let mut writer = AsyncArrowWriter::try_new(store_writer, schema.clone(), Some(props))?; - - // Data sketches quantiles sketches - let mut sketches = make_sketches_for_schema(schema, &row_key_fields); - - // Container for errors - let errors: RefCell> = RefCell::new(vec![]); - - let error_free = - iters.map(|iter| iter.filter_map(|r| r.map_err(|e| errors.borrow_mut().push(e)).ok())); - - // Row K-way merge on iterators - let merged = kmerge(error_free); - - // Need to collect min/max key on column 0. Not sure we can assume col. 0 is a row key column, so it might not have a quantiles sketch - let mut rows_written = 0; - for chunk in &merged.chunks(row_group_size) { - // Check for errors - if !errors.borrow().is_empty() { - // Return first error - return Err(errors.borrow_mut().remove(0)); - } - - let rows = chunk.collect::>(); - let rows = rows - .iter() - .map(arrow::row::OwnedRow::row) - .collect::>(); - - rows_written += rows.len(); - info!( - "Merged {} rows", - rows_written.to_formatted_string(&Locale::en) - ); - - let cols = converter_ptr.borrow_mut().convert_rows(rows)?; - let batch = RecordBatch::try_new(schema.clone(), cols)?; - - // update the data sketches on this batch - update_sketches(&batch, &mut sketches, &row_key_fields); - - // write some data sketches here - batch.column(0); - - writer.write(&batch).await?; } - - writer.close().await?; - - // serialise the sketches - let sketch_path = create_sketch_path(output_file_path); - serialise_sketches(store_factory, &sketch_path, &sketches)?; - - info!( - "Object store made {} GET requests and read {} bytes", - store - .get_count() - .unwrap_or_default() - .to_formatted_string(&Locale::en), - store - .get_bytes_read() - .unwrap_or_default() - .to_formatted_string(&Locale::en), - ); - - Ok(CompactionResult { - rows_read: rows_written, - rows_written, - }) } /// Creates a file path suitable for writing sketches to. @@ -465,157 +165,3 @@ pub fn create_sketch_path(output_path: &Url) -> Url { ); res } - -/// Creates an iterator of Arrow record batches. -/// -/// This creates a file reader for a Parquet file returning the Arrow recordbatches. -/// The optional `batch_size` argument allows for setting the maximum size of each batch -/// returned in one go. -/// -/// # Errors -/// The supplied file path must have a valid extension: parquet, parq, or pq. -pub async fn get_file_iterator( - store_factory: &ObjectStoreFactory, - file_path: &Url, - batch_size: Option, -) -> Result>, ArrowError> { - get_file_iterator_for_row_group_range(store_factory, file_path, None, None, batch_size).await -} - -/// Creates an iterator of Arrow record batches from a Parquet file. -/// -/// This allows for the row groups to be read from the Parquet file -/// to be specified. -/// -/// # Examples -/// ```ignore -/// # use url::Url; -/// -/// let s = get_file_iterator_for_row_group_range(store_factory, Url::parse("/path/to/file.parquet").unwrap(), Some(0), Some(10), None); -/// ``` -/// -/// # Errors -/// The supplied file path must have a valid extension: parquet, parq, or pq. -/// -/// The end row group, if specified, must be greater than or equal to the start row group. -async fn get_file_iterator_for_row_group_range( - store_factory: &ObjectStoreFactory, - src: &Url, - start_row_group: Option, - end_row_group: Option, - batch_size: Option, -) -> Result>, ArrowError> { - let extension = PathBuf::from(src.path()) - .extension() - .map(|ext| ext.to_str().unwrap_or_default().to_owned()); - if let Some("parquet" | "parq" | "pq") = extension.as_deref() { - let mut builder = get_parquet_builder(store_factory, src).await?; - - if let Some(num_groups) = batch_size { - builder = builder.with_batch_size(num_groups); - } - - let num_row_groups = builder.metadata().num_row_groups(); - let num_rows = builder.metadata().file_metadata().num_rows(); - - let rg_start = start_row_group.unwrap_or(0); - let rg_end = end_row_group.unwrap_or(num_row_groups); - - if rg_end < rg_start { - return Err(ArrowError::InvalidArgumentError( - "end row group must not be less than start".to_owned(), - )); - } - - info!( - "Creating input iterator for file {} row groups {} rows {} loading row groups {}..{}", - src.as_ref(), - num_row_groups.to_formatted_string(&Locale::en), - num_rows.to_formatted_string(&Locale::en), - rg_start.to_formatted_string(&Locale::en), - rg_end.to_formatted_string(&Locale::en) - ); - - builder = builder.with_row_groups((rg_start..rg_end).collect()); - - builder - .build() - .map(futures::executor::block_on_stream) - .map_err(ArrowError::from) - } else { - let p = extension.unwrap_or("".to_owned()); - Err(ArrowError::InvalidArgumentError(format!( - "Unrecognised extension {p}" - ))) - } -} - -/// Create an asynchronous builder for reading Parquet files from an object store. -/// -/// The URL must start with a scheme that the object store recognises, e.g. "file" or "s3". -/// -/// # Errors -/// This function will return an error if it couldn't connect to S3 or open a valid -/// Parquet file. -pub async fn get_parquet_builder( - store_factory: &ObjectStoreFactory, - src: &Url, -) -> Result, ArrowError> { - let store = store_factory.get_object_store(src)?; - // HEAD the file to get metadata - let path = object_store::path::Path::from(src.path()); - let object_meta = store - .head(&path) - .await - .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; - // Create a reader for the target file, use it to construct a Stream - let reader = ParquetObjectReader::new(store.as_object_store(), object_meta); - Ok(ParquetRecordBatchStreamBuilder::new(reader).await?) -} - -/// Read the schema from the given Parquet file. -/// -/// The Parquet file metadata is read and returned. -/// -/// # Examples -/// ```no_run -/// # use compaction::read_schema; -/// # use std::sync::Arc; -/// # use arrow::datatypes::Schema; -/// # use aws_types::region::Region; -/// # use url::Url; -/// # use crate::compaction::ObjectStoreFactory; -/// let sf = ObjectStoreFactory::new(None, &Region::new("eu-west-2")); -/// let schema: Arc = read_schema(&sf, &Url::parse("file:///path/to/my/file").unwrap()).unwrap(); -/// ``` -/// # Errors -/// Errors can occur if we are not able to read the named file or if the schema/file is corrupt. -/// -pub async fn read_schema( - store_factory: &ObjectStoreFactory, - src: &Url, -) -> Result, ArrowError> { - let builder = get_parquet_builder(store_factory, src).await?; - Ok(builder.schema().clone()) -} - -/// Checks all files match the schema. -/// -/// The schema should be provided from the first file, all subsequent -/// file schemas are compared to this. This function returns true if all match. -#[must_use] -pub async fn validate_schemas_same( - store_factory: &ObjectStoreFactory, - input_file_paths: &[Url], - schema: &Arc, -) -> bool { - for path in input_file_paths { - let file_schema = read_schema(store_factory, path) - .await - .unwrap_or(Arc::new(Schema::empty())); - if file_schema != *schema { - return false; - } - } - true -} From b1f76ddaadcbcf10875bfa674ae54e08939c68eb Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 23 Apr 2024 15:26:51 +0000 Subject: [PATCH 045/129] License changes --- rust/Cargo.lock | 604 +++++++++++++++++- rust/Cargo.toml | 2 +- rust/Cross.toml | 13 + rust/compaction/.cargo/config.toml | 2 - rust/compaction/Cargo.toml | 11 +- rust/compaction/src/aws_s3.rs | 2 +- rust/compaction/src/datafusion.rs | 19 + rust/compaction/src/details.rs | 28 +- rust/compaction/src/lib.rs | 7 +- rust/compaction/src/sketch.rs | 2 +- rust/compactor/Cargo.toml | 2 +- rust/compactor/src/bin/main.rs | 2 +- .../src/include/byte_array_serializer.hpp | 4 +- rust/rust_sketch/src/include/common_types.hpp | 8 +- rust/rust_sketch/src/include/quantiles.hpp | 10 +- rust/rust_sketch/src/lib.rs | 2 +- rust/rust_sketch/src/main.rs | 2 +- rust/rust_sketch/src/quantiles.rs | 2 +- 18 files changed, 661 insertions(+), 61 deletions(-) delete mode 100644 rust/compaction/.cargo/config.toml create mode 100644 rust/compaction/src/datafusion.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 86180d9fd0..501014fc6a 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -55,6 +55,12 @@ dependencies = [ "alloc-no-stdlib", ] +[[package]] +name = "allocator-api2" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f" + [[package]] name = "android-tzdata" version = "0.1.1" @@ -119,10 +125,10 @@ dependencies = [ ] [[package]] -name = "anyhow" -version = "1.0.82" +name = "arrayref" +version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f538837af36e6f6a9be0faa67f9a314f8119e4e4b5867c6ab40ed60360142519" +checksum = "6b4930d2cb77ce62f89ee5d5289b4ac049559b1c45539271f5ed4fdc7db34545" [[package]] name = "arrayvec" @@ -177,6 +183,7 @@ dependencies = [ "arrow-data", "arrow-schema", "chrono", + "chrono-tz", "half", "hashbrown", "num", @@ -207,6 +214,7 @@ dependencies = [ "atoi", "base64 0.22.0", "chrono", + "comfy-table", "half", "lexical-core", "num", @@ -256,6 +264,7 @@ dependencies = [ "arrow-data", "arrow-schema", "flatbuffers", + "lz4_flex", ] [[package]] @@ -345,6 +354,24 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "async-compression" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07dbbf24db18d609b1462965249abdf49129ccad073ec257da372adc83259c60" +dependencies = [ + "bzip2", + "flate2", + "futures-core", + "futures-io", + "memchr", + "pin-project-lite", + "tokio", + "xz2", + "zstd", + "zstd-safe", +] + [[package]] name = "async-stream" version = "0.3.5" @@ -738,6 +765,28 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf4b9d6a944f767f8e5e0db018570623c85f3d925ac718db4e06d0187adb21c1" +[[package]] +name = "blake2" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" +dependencies = [ + "digest", +] + +[[package]] +name = "blake3" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30cca6d3674597c30ddf2c587bf8d9d65c9a84d2326d941cc79c9842dfe0ef52" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -796,6 +845,27 @@ dependencies = [ "either", ] +[[package]] +name = "bzip2" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8" +dependencies = [ + "bzip2-sys", + "libc", +] + +[[package]] +name = "bzip2-sys" +version = "0.1.11+1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "736a955f3fa7875102d57c82b8cac37ec45224a07fd32d58f9f7a186b6cd4cdc" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "cc" version = "1.0.95" @@ -828,6 +898,28 @@ dependencies = [ "windows-targets 0.52.5", ] +[[package]] +name = "chrono-tz" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d59ae0466b83e838b81a54256c39d5d7c20b9d7daa10510a242d9b75abd5936e" +dependencies = [ + "chrono", + "chrono-tz-build", + "phf", +] + +[[package]] +name = "chrono-tz-build" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "433e39f13c9a060046954e0592a8d0a4bcb1040125cbf91cb8ee58964cfb350f" +dependencies = [ + "parse-zoneinfo", + "phf", + "phf_codegen", +] + [[package]] name = "clap" version = "4.5.4" @@ -911,26 +1003,36 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" +[[package]] +name = "comfy-table" +version = "7.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b34115915337defe99b2aff5c2ce6771e5fbc4079f4b506301f5cf394c8452f7" +dependencies = [ + "strum", + "strum_macros", + "unicode-width", +] + [[package]] name = "compaction" version = "0.1.0" dependencies = [ - "anyhow", "arrow", "aws-config", "aws-credential-types", "aws-types", "bytes", "chrono", + "color-eyre", "cxx", + "datafusion", "env_logger", "futures", - "itertools", "libc", "log", "num-format", "object_store", - "parquet", "rust_sketch", "tokio", "tokio-test", @@ -977,6 +1079,12 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "constant_time_eq" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7144d30dcf0fafbce74250a3963025d8d52177934239851c917d29f1df280c2" + [[package]] name = "core-foundation" version = "0.9.4" @@ -1092,6 +1200,283 @@ dependencies = [ "syn 2.0.60", ] +[[package]] +name = "dashmap" +version = "5.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +dependencies = [ + "cfg-if", + "hashbrown", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "datafusion" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85069782056753459dc47e386219aa1fdac5b731f26c28abb8c0ffd4b7c5ab11" +dependencies = [ + "ahash", + "arrow", + "arrow-array", + "arrow-ipc", + "arrow-schema", + "async-compression", + "async-trait", + "bytes", + "bzip2", + "chrono", + "dashmap", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions", + "datafusion-functions-array", + "datafusion-optimizer", + "datafusion-physical-expr", + "datafusion-physical-plan", + "datafusion-sql", + "flate2", + "futures", + "glob", + "half", + "hashbrown", + "indexmap", + "itertools", + "log", + "num_cpus", + "object_store", + "parking_lot", + "parquet", + "pin-project-lite", + "rand", + "sqlparser", + "tempfile", + "tokio", + "tokio-util", + "url", + "uuid", + "xz2", + "zstd", +] + +[[package]] +name = "datafusion-common" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "309d9040751f6dc9e33c85dce6abb55a46ef7ea3644577dd014611c379447ef3" +dependencies = [ + "ahash", + "arrow", + "arrow-array", + "arrow-buffer", + "arrow-schema", + "chrono", + "half", + "instant", + "libc", + "num_cpus", + "object_store", + "parquet", + "sqlparser", +] + +[[package]] +name = "datafusion-common-runtime" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3e4a44d8ef1b1e85d32234e6012364c411c3787859bb3bba893b0332cb03dfd" +dependencies = [ + "tokio", +] + +[[package]] +name = "datafusion-execution" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06a3a29ae36bcde07d179cc33b45656a8e7e4d023623e320e48dcf1200eeee95" +dependencies = [ + "arrow", + "chrono", + "dashmap", + "datafusion-common", + "datafusion-expr", + "futures", + "hashbrown", + "log", + "object_store", + "parking_lot", + "rand", + "tempfile", + "url", +] + +[[package]] +name = "datafusion-expr" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a3542aa322029c2121a671ce08000d4b274171070df13f697b14169ccf4f628" +dependencies = [ + "ahash", + "arrow", + "arrow-array", + "chrono", + "datafusion-common", + "paste", + "sqlparser", + "strum", + "strum_macros", +] + +[[package]] +name = "datafusion-functions" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd221792c666eac174ecc09e606312844772acc12cbec61a420c2fca1ee70959" +dependencies = [ + "arrow", + "base64 0.22.0", + "blake2", + "blake3", + "chrono", + "datafusion-common", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "hex", + "itertools", + "log", + "md-5", + "regex", + "sha2", + "unicode-segmentation", + "uuid", +] + +[[package]] +name = "datafusion-functions-array" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e501801e84d9c6ef54caaebcda1b18a6196a24176c12fb70e969bc0572e03c55" +dependencies = [ + "arrow", + "arrow-array", + "arrow-buffer", + "arrow-ord", + "arrow-schema", + "datafusion-common", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions", + "itertools", + "log", + "paste", +] + +[[package]] +name = "datafusion-optimizer" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76bd7f5087817deb961764e8c973d243b54f8572db414a8f0a8f33a48f991e0a" +dependencies = [ + "arrow", + "async-trait", + "chrono", + "datafusion-common", + "datafusion-expr", + "datafusion-physical-expr", + "hashbrown", + "itertools", + "log", + "regex-syntax", +] + +[[package]] +name = "datafusion-physical-expr" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cabc0d9aaa0f5eb1b472112f16223c9ffd2fb04e58cbf65c0a331ee6e993f96" +dependencies = [ + "ahash", + "arrow", + "arrow-array", + "arrow-buffer", + "arrow-ord", + "arrow-schema", + "arrow-string", + "base64 0.22.0", + "blake2", + "blake3", + "chrono", + "datafusion-common", + "datafusion-execution", + "datafusion-expr", + "half", + "hashbrown", + "hex", + "indexmap", + "itertools", + "log", + "md-5", + "paste", + "petgraph", + "rand", + "regex", + "sha2", + "unicode-segmentation", +] + +[[package]] +name = "datafusion-physical-plan" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17c0523e9c8880f2492a88bbd857dde02bed1ed23f3e9211a89d3d7ec3b44af9" +dependencies = [ + "ahash", + "arrow", + "arrow-array", + "arrow-buffer", + "arrow-schema", + "async-trait", + "chrono", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "futures", + "half", + "hashbrown", + "indexmap", + "itertools", + "log", + "once_cell", + "parking_lot", + "pin-project-lite", + "rand", + "tokio", +] + +[[package]] +name = "datafusion-sql" +version = "37.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49eb54b42227136f6287573f2434b1de249fe1b8e6cd6cc73a634e4a3ec29356" +dependencies = [ + "arrow", + "arrow-array", + "arrow-schema", + "datafusion-common", + "datafusion-expr", + "log", + "sqlparser", + "strum", +] + [[package]] name = "deranged" version = "0.3.11" @@ -1162,6 +1547,16 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +[[package]] +name = "errno" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a258e46cdc063eb8519c00b9fc845fc47bcfca4130e2f08e88665ceda8474245" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "eyre" version = "0.6.12" @@ -1178,6 +1573,12 @@ version = "2.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "658bd65b1cf4c852a3cc96f18a8ce7b5640f6b703f905c7d74532294c2a63984" +[[package]] +name = "fixedbitset" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" + [[package]] name = "flatbuffers" version = "23.5.26" @@ -1359,6 +1760,12 @@ dependencies = [ "url", ] +[[package]] +name = "glob" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" + [[package]] name = "h2" version = "0.3.26" @@ -1394,6 +1801,10 @@ name = "hashbrown" version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "290f1a1d9242c78d09ce40a5e87e7554ee637af1351968159f4952f028f75604" +dependencies = [ + "ahash", + "allocator-api2", +] [[package]] name = "heck" @@ -1607,6 +2018,18 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "instant" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c" +dependencies = [ + "cfg-if", + "js-sys", + "wasm-bindgen", + "web-sys", +] + [[package]] name = "integer-encoding" version = "3.0.4" @@ -1783,6 +2206,12 @@ dependencies = [ "cc", ] +[[package]] +name = "linux-raw-sys" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01cda141df6706de531b6c46c3a33ecca755538219bd484262fa09410c13539c" + [[package]] name = "lock_api" version = "0.4.11" @@ -1808,6 +2237,17 @@ dependencies = [ "twox-hash", ] +[[package]] +name = "lzma-sys" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "md-5" version = "0.10.6" @@ -2147,6 +2587,15 @@ dependencies = [ "zstd", ] +[[package]] +name = "parse-zoneinfo" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c705f256449c60da65e11ff6626e0c16a0a0b96aaa348de61376b249bc340f41" +dependencies = [ + "regex", +] + [[package]] name = "paste" version = "1.0.14" @@ -2159,6 +2608,54 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[package]] +name = "petgraph" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d3afd2628e69da2be385eb6f2fd57c8ac7977ceeff6dc166ff1657b0e386a9" +dependencies = [ + "fixedbitset", + "indexmap", +] + +[[package]] +name = "phf" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_codegen" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8d39688d359e6b34654d328e262234662d16cc0f60ec8dcbe5e718709342a5a" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48e4cc64c2ad9ebe670cb8fd69dd50ae301650392e81c05f9bfcb2d5bdbc24b0" +dependencies = [ + "phf_shared", + "rand", +] + +[[package]] +name = "phf_shared" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project-lite" version = "0.2.14" @@ -2375,6 +2872,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "0.38.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70dc5ec042f7a43c4a73241207cecc9873a06d45debb38b329f8541d85c2730f" +dependencies = [ + "bitflags 2.5.0", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.52.0", +] + [[package]] name = "rustls" version = "0.21.11" @@ -2418,6 +2928,12 @@ dependencies = [ "untrusted", ] +[[package]] +name = "rustversion" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80af6f9131f277a45a3fba6ce8e2258037bb0477a67e610d3c1fe046ab31de47" + [[package]] name = "ryu" version = "1.0.17" @@ -2580,6 +3096,12 @@ dependencies = [ "libc", ] +[[package]] +name = "siphasher" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" + [[package]] name = "slab" version = "0.4.9" @@ -2639,6 +3161,27 @@ version = "0.9.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +[[package]] +name = "sqlparser" +version = "0.44.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aaf9c7ff146298ffda83a200f8d5084f08dcee1edfc135fcc1d646a45d50ffd6" +dependencies = [ + "log", + "sqlparser_derive", +] + +[[package]] +name = "sqlparser_derive" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01b2e185515564f15375f593fb966b5718bc624ba77fe49fa4616ad619690554" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.60", +] + [[package]] name = "static_assertions" version = "1.1.0" @@ -2651,6 +3194,28 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d8cec3501a5194c432b2b7976db6b7d10ec95c253208b45f83f7136aa985e29" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6cf59daf282c0a494ba14fd21610a0325f9f90ec9d1231dea26bcb1d696c946" +dependencies = [ + "heck 0.4.1", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.60", +] + [[package]] name = "subtle" version = "2.5.0" @@ -2706,6 +3271,18 @@ dependencies = [ "libc", ] +[[package]] +name = "tempfile" +version = "3.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85b77fafb263dd9d05cbeac119526425676db3784113aa9295c88498cbf8bff1" +dependencies = [ + "cfg-if", + "fastrand", + "rustix", + "windows-sys 0.52.0", +] + [[package]] name = "termcolor" version = "1.4.1" @@ -3023,6 +3600,12 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "unicode-segmentation" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4c87d22b6e3f4a18d4d40ef354e97c90fcb14dd91d7dc0aa9d8a1172ebf7202" + [[package]] name = "unicode-width" version = "0.1.11" @@ -3400,6 +3983,15 @@ version = "0.13.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" +[[package]] +name = "xz2" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" +dependencies = [ + "lzma-sys", +] + [[package]] name = "zerocopy" version = "0.7.32" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index b1c13d29f8..1b0a27d719 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,4 +1,4 @@ -# Copyright 2022-2023 Crown Copyright +# 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. diff --git a/rust/Cross.toml b/rust/Cross.toml index d2d535f49a..508ac28b39 100644 --- a/rust/Cross.toml +++ b/rust/Cross.toml @@ -1,3 +1,16 @@ +# 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. [target.aarch64-unknown-linux-gnu] image = "ghcr.io/cross-rs/aarch64-unknown-linux-gnu:edge" diff --git a/rust/compaction/.cargo/config.toml b/rust/compaction/.cargo/config.toml deleted file mode 100644 index 656e08b011..0000000000 --- a/rust/compaction/.cargo/config.toml +++ /dev/null @@ -1,2 +0,0 @@ -[net] -git-fetch-with-cli = true \ No newline at end of file diff --git a/rust/compaction/Cargo.toml b/rust/compaction/Cargo.toml index 30fdab39df..f0a7392895 100644 --- a/rust/compaction/Cargo.toml +++ b/rust/compaction/Cargo.toml @@ -1,4 +1,4 @@ -# Copyright 2022-2023 Crown Copyright +# 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. @@ -35,14 +35,9 @@ crate-type = ["cdylib", "rlib"] cxx = { version = "1.0.104" } # Exception handling for Rust libc = { version = "0.2.147" } # FFI type support log = { version = "0.4.19" } # Logging support -itertools = { version = "0.12.1" } # kmerge arrow = { version = "51.0.0" } # Batch of reading from Parquet files futures = { version = "0.3.28" } # Async processing -parquet = { version = "51.0.0", features = [ - "arrow", - "async", - "object_store", -] } # Parquet processing +datafusion = { version = "37.1.0", features = ["backtrace"] } object_store = { version = "0.9.1", features = [ "aws", ] } # Remote cloud storage access @@ -57,4 +52,4 @@ num-format = { version = "0.4.4" } # Nicely formatted numbers tokio-test = { version = "0.4.2" } # Doc tests env_logger = { version = "0.11.3" } chrono = { version = "0.4.26" } # Log helper -anyhow = { version = "1.0.82" } +color-eyre = { version = "0.6.3" } # Error handling diff --git a/rust/compaction/src/aws_s3.rs b/rust/compaction/src/aws_s3.rs index 4d0c514821..af12f7bc24 100644 --- a/rust/compaction/src/aws_s3.rs +++ b/rust/compaction/src/aws_s3.rs @@ -2,7 +2,7 @@ //! //! This module contains support functions and structs for accessing AWS S3 via the [`object_store`] crate. /* - * Copyright 2022-2023 Crown Copyright + * 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. diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs new file mode 100644 index 0000000000..dc395cfdf9 --- /dev/null +++ b/rust/compaction/src/datafusion.rs @@ -0,0 +1,19 @@ +/// `DataFusion` contains the implementation for performing Sleeper compactions +/// using Apache DataFusion. +/// +/// This allows for multi-threaded compaction and optimised Parquet reading. +/* + * 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. + */ \ No newline at end of file diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 1971e540bc..13582c448d 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -15,32 +15,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -use crate::{ - aws_s3::ObjectStoreFactory, - sketch::{make_sketches_for_schema, serialise_sketches, update_sketches}, -}; -use arrow::{ - datatypes::Schema, - error::ArrowError, - record_batch::RecordBatch, - row::{OwnedRow, RowConverter, Rows, SortField}, -}; +use crate::aws_s3::ObjectStoreFactory; +use arrow::error::ArrowError; use aws_credential_types::Credentials; use aws_types::region::Region; -use futures::{executor::BlockingStream, future}; -use itertools::{kmerge, Itertools}; -use log::{debug, info}; -use num_format::{Locale, ToFormattedString}; -use parquet::{ - arrow::{ - async_reader::{ParquetObjectReader, ParquetRecordBatchStream}, - AsyncArrowWriter, ParquetRecordBatchStreamBuilder, - }, - basic::{Compression, ZstdLevel}, - file::properties::WriterProperties, -}; -use std::{cell::RefCell, collections::HashMap, path::PathBuf, sync::Arc}; +use std::{collections::HashMap, path::PathBuf}; use url::Url; /// Type safe variant for Sleeper partition boundary @@ -146,6 +126,8 @@ pub async fn merge_sorted_files( if output_file_path.scheme() == "s3a" { let _ = output_file_path.set_scheme("s3"); } + + Ok(CompactionResult { rows_read: 0, rows_written: 0, diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index b38454345b..b8260ef4c6 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -6,7 +6,7 @@ //! We have an internal "details" module that encapsulates the internal workings. All the //! public API should be in this module. /* - * Copyright 2022-2023 Crown Copyright + * 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. @@ -22,6 +22,7 @@ */ mod aws_s3; +mod datafusion; mod details; mod sketch; @@ -132,7 +133,7 @@ pub struct FFICompactionParams { } impl TryFrom<&FFICompactionParams> for CompactionInput { - type Error = anyhow::Error; + type Error = color_eyre::eyre::Report; fn try_from(params: &FFICompactionParams) -> Result { // We do this separately since we need the values for computing the region @@ -176,7 +177,7 @@ impl TryFrom<&FFICompactionParams> for CompactionInput { fn compute_region>( params: &FFICompactionParams, row_key_cols: &[T], -) -> Result, anyhow::Error> { +) -> color_eyre::Result> { let region_mins_inclusive = unpack_primitive_array( params.region_mins_inclusive, params.region_mins_inclusive_len, diff --git a/rust/compaction/src/sketch.rs b/rust/compaction/src/sketch.rs index 641aae773d..9bb6512b70 100644 --- a/rust/compaction/src/sketch.rs +++ b/rust/compaction/src/sketch.rs @@ -1,7 +1,7 @@ //! This module provides some extra required wrappers around sketches //! functionality such as common traits. /* - * Copyright 2022-2023 Crown Copyright + * 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. diff --git a/rust/compactor/Cargo.toml b/rust/compactor/Cargo.toml index f647000c9d..2367910848 100644 --- a/rust/compactor/Cargo.toml +++ b/rust/compactor/Cargo.toml @@ -1,4 +1,4 @@ -# Copyright 2022-2023 Crown Copyright +# 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. diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index c098b61538..14690903f1 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -1,5 +1,5 @@ /* - * Copyright 2022-2023 Crown Copyright + * 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. diff --git a/rust/rust_sketch/src/include/byte_array_serializer.hpp b/rust/rust_sketch/src/include/byte_array_serializer.hpp index be2f8be01f..a08c348b41 100644 --- a/rust/rust_sketch/src/include/byte_array_serializer.hpp +++ b/rust/rust_sketch/src/include/byte_array_serializer.hpp @@ -1,5 +1,5 @@ /* - * Copyright 2022-2023 Crown Copyright + * 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. @@ -19,7 +19,7 @@ * @brief A serializer for byte array sketches. This was adapted from datasketches::serde * @date 2023 * - * @copyright Copyright 2022-2023 Crown Copyright + * @copyright Copyright 2022-2024 Crown Copyright * */ #include diff --git a/rust/rust_sketch/src/include/common_types.hpp b/rust/rust_sketch/src/include/common_types.hpp index e8bdbf795a..0a5ad83daa 100644 --- a/rust/rust_sketch/src/include/common_types.hpp +++ b/rust/rust_sketch/src/include/common_types.hpp @@ -1,5 +1,5 @@ /* - * Copyright 2022-2023 Crown Copyright + * 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. @@ -18,9 +18,9 @@ * @file common_types.hpp * @brief Common type declarations. * @date 2023 - * - * @copyright Copyright 2022-2023 Crown Copyright - * + * + * @copyright Copyright 2022-2024 Crown Copyright + * */ #include #include diff --git a/rust/rust_sketch/src/include/quantiles.hpp b/rust/rust_sketch/src/include/quantiles.hpp index 6af355d981..c6a8c4f3ae 100644 --- a/rust/rust_sketch/src/include/quantiles.hpp +++ b/rust/rust_sketch/src/include/quantiles.hpp @@ -1,5 +1,5 @@ /* - * Copyright 2022-2023 Crown Copyright + * 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. @@ -18,9 +18,9 @@ * @file quantiles.hpp * @brief Contains the wrapping and helper functions to facilitate mapping from Rust to C++ * for quantiles sketches. - * @date 2023 + * @date 2024 * - * @copyright Copyright 2022-2023 Crown Copyright + * @copyright Copyright 2022-2024 Crown Copyright * */ @@ -44,7 +44,7 @@ namespace rust_sketch { /** * @brief Use our custom byte_array serializer for byte_array type, otherwise fallback to default one provided by datasketches library. - * @tparam T element type of sketch + * @tparam T element type of sketch */ template using serializer_t = std::conditional_t, ByteArraySerde, datasketches::serde>; @@ -207,7 +207,7 @@ struct quantiles_sketch_derived : public base_type { /** * @brief Merges another sketch into this one. - * + * * @param other the other sketch to merge into this one */ template diff --git a/rust/rust_sketch/src/lib.rs b/rust/rust_sketch/src/lib.rs index 8b22b4c009..09a7447db8 100644 --- a/rust/rust_sketch/src/lib.rs +++ b/rust/rust_sketch/src/lib.rs @@ -36,7 +36,7 @@ //! wrapped to return a [`Result`]. //! /* - * Copyright 2022-2023 Crown Copyright + * 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. diff --git a/rust/rust_sketch/src/main.rs b/rust/rust_sketch/src/main.rs index 8251265593..78e918fec9 100644 --- a/rust/rust_sketch/src/main.rs +++ b/rust/rust_sketch/src/main.rs @@ -1,5 +1,5 @@ /* - * Copyright 2022-2023 Crown Copyright + * 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. diff --git a/rust/rust_sketch/src/quantiles.rs b/rust/rust_sketch/src/quantiles.rs index 003d9f3f4a..56c64634c6 100644 --- a/rust/rust_sketch/src/quantiles.rs +++ b/rust/rust_sketch/src/quantiles.rs @@ -54,7 +54,7 @@ //! ); //! ``` /* - * Copyright 2022-2023 Crown Copyright + * 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. From 33ce8951284ba06d2af6f1f00ab68533c236476e Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 23 Apr 2024 16:04:21 +0000 Subject: [PATCH 046/129] DataFusion skeleton ready --- rust/Cargo.lock | 7 ++--- rust/compaction/src/datafusion.rs | 47 ++++++++++++++++++++----------- rust/compaction/src/details.rs | 46 +++++++++++++++++------------- rust/compaction/src/lib.rs | 46 +++++++++--------------------- rust/compactor/Cargo.toml | 5 +--- rust/compactor/src/bin/main.rs | 14 +-------- 6 files changed, 74 insertions(+), 91 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 501014fc6a..af0f6504be 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1043,9 +1043,6 @@ dependencies = [ name = "compactor" version = "0.1.0" dependencies = [ - "aws-config", - "aws-credential-types", - "aws-types", "chrono", "clap", "color-eyre", @@ -1909,9 +1906,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "human-panic" -version = "1.2.3" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4f016c89920bbb30951a8405ecacbb4540db5524313b9445736e7e1855cf370" +checksum = "a4c5d0e9120f6bca6120d142c7ede1ba376dd6bf276d69dd3dbe6cbeb7824179" dependencies = [ "anstream", "anstyle", diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index dc395cfdf9..3e017f5876 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -1,19 +1,34 @@ /// `DataFusion` contains the implementation for performing Sleeper compactions -/// using Apache DataFusion. -/// +/// using Apache `DataFusion`. +/// /// This allows for multi-threaded compaction and optimised Parquet reading. /* - * 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. - */ \ No newline at end of file +* 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. +*/ +use crate::{aws_s3::ObjectStoreFactory, CompactionInput, CompactionResult}; +use datafusion::error::DataFusionError; +use url::Url; + +pub async fn compact( + store_factory: &ObjectStoreFactory, + input_data: &CompactionInput, + input_paths: &[Url], + output_path: &Url, +) -> Result { + Ok(CompactionResult { + rows_read: 0, + rows_written: 0, + }) +} diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 13582c448d..be85b8d4d9 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -16,9 +16,9 @@ * limitations under the License. */ use crate::aws_s3::ObjectStoreFactory; -use arrow::error::ArrowError; -use aws_credential_types::Credentials; -use aws_types::region::Region; +use aws_config::BehaviorVersion; +use aws_credential_types::provider::ProvideCredentials; +use color_eyre::eyre::{eyre, Result}; use std::{collections::HashMap, path::PathBuf}; use url::Url; @@ -93,21 +93,11 @@ pub struct CompactionResult { /// # Errors /// There must be at least one input file. /// -#[allow(clippy::too_many_arguments)] -#[allow(clippy::arc_with_non_send_sync)] -pub async fn merge_sorted_files( - aws_creds: Option, - region: &Region, - input_data: &CompactionInput, -) -> Result { +pub async fn merge_sorted_files(input_data: &CompactionInput) -> Result { // Read the schema from the first file if input_data.input_files.is_empty() { - Err(ArrowError::InvalidArgumentError( - "No input paths supplied".into(), - )) + Err(eyre!("No input paths supplied")) } else { - // Create our object store factory - let store_factory = ObjectStoreFactory::new(aws_creds, region); // Java tends to use s3a:// URI scheme instead of s3:// so map it here let input_file_paths: Vec = input_data .input_files @@ -127,11 +117,27 @@ pub async fn merge_sorted_files( let _ = output_file_path.set_scheme("s3"); } - - Ok(CompactionResult { - rows_read: 0, - rows_written: 0, - }) + let config = aws_config::defaults(BehaviorVersion::latest()).load().await; + let region = config + .region() + .ok_or(eyre!("Couldn't retrieve AWS region"))?; + let creds: aws_credential_types::Credentials = config + .credentials_provider() + .ok_or(eyre!("Couldn't retrieve AWS credentials"))? + .provide_credentials() + .await?; + + // Create our object store factory + let store_factory = ObjectStoreFactory::new(Some(creds), region); + + crate::datafusion::compact( + &store_factory, + input_data, + &input_file_paths, + &output_file_path, + ) + .await + .map_err(Into::into) } } diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index b8260ef4c6..76fe58152e 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -26,13 +26,10 @@ mod datafusion; mod details; mod sketch; -use arrow::error::ArrowError; -use aws_config::BehaviorVersion; -use aws_credential_types::provider::ProvideCredentials; use chrono::Local; -use futures::TryFutureExt; +use color_eyre::eyre::eyre; use libc::{c_void, size_t, EFAULT, EINVAL, EIO}; -use log::{error, info, LevelFilter}; +use log::{error, info, warn, LevelFilter}; use std::borrow::Borrow; use std::collections::HashMap; use std::io::Write; @@ -77,27 +74,6 @@ fn maybe_cfg_log() { }); } -/// Obtains AWS credentials from normal places and then calls merge function -/// with obtained credentials. -/// -async fn credentials_and_merge( - input_data: &CompactionInput, -) -> Result { - let config = aws_config::defaults(BehaviorVersion::latest()).load().await; - let region = config.region().ok_or(ArrowError::InvalidArgumentError( - "Couldn't retrieve AWS region".into(), - ))?; - let creds: aws_credential_types::Credentials = config - .credentials_provider() - .ok_or(ArrowError::InvalidArgumentError( - "Couldn't retrieve AWS credentials".into(), - ))? - .provide_credentials() - .map_err(|e| ArrowError::ExternalError(Box::new(e))) - .await?; - merge_sorted_files(Some(creds), region, input_data).await -} - /// Contains all the input data for setting up a compaction. /// /// See java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java @@ -271,13 +247,19 @@ pub extern "C" fn ffi_merge_sorted_files( output_data: *mut FFICompactionResult, ) -> c_int { maybe_cfg_log(); + if let Err(e) = color_eyre::install() { + warn!("Couldn't install color_eyre error handler"); + } let Some(params) = (unsafe { input_data.as_ref() }) else { error!("input data pointer is null"); return EFAULT; }; // Start async runtime - let rt = match tokio::runtime::Runtime::new() { + let rt = match tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + { Ok(v) => v, Err(e) => { error!("Couldn't create Rust tokio runtime {}", e); @@ -294,7 +276,7 @@ pub extern "C" fn ffi_merge_sorted_files( }; // Run compaction - let result = rt.block_on(credentials_and_merge(&details)); + let result = rt.block_on(merge_sorted_files(&details)); match result { Ok(res) => { if let Some(data) = unsafe { output_data.as_mut() } { @@ -320,7 +302,7 @@ pub extern "C" fn ffi_merge_sorted_files( fn unpack_string_array( array_base: *const *const c_char, len: usize, -) -> Result, ArrowError> { +) -> color_eyre::Result> { unsafe { // create a slice from the pointer slice::from_raw_parts(array_base, len) @@ -338,16 +320,14 @@ fn unpack_string_array( #[allow(clippy::cast_ptr_alignment)] let str_len = unsafe { *(*s).cast::() }; if str_len < 0 { - return Err(ArrowError::InvalidArgumentError(format!( - "Illegal string length in FFI array: {str_len}" - ))); + return Err(eyre!("Illegal string length in FFI array: {str_len}")); } // convert to string and check it's valid std::str::from_utf8(unsafe { #[allow(clippy::cast_sign_loss)] slice::from_raw_parts(s.byte_add(4).cast::(), str_len as usize) }) - .map_err(|e| ArrowError::ExternalError(Box::new(e))) + .map_err(Into::into) }) // now convert to a vector if all strings OK, else Err .collect::, _>>() diff --git a/rust/compactor/Cargo.toml b/rust/compactor/Cargo.toml index 2367910848..447792e603 100644 --- a/rust/compactor/Cargo.toml +++ b/rust/compactor/Cargo.toml @@ -32,14 +32,11 @@ log = { version = "0.4.17", features = [ "release_max_level_debug", ] } # Standard logging framework env_logger = { version = "0.11.3" } # Standard logging to stderr -human-panic = { version = "1.1.4" } # Readable panic messages +human-panic = { version = "2.0.0" } # Readable panic messages clap = { version = "4.3.0", features = ["derive"] } # Cmd line args processing color-eyre = { version = "0.6.2" } # Colourised version of `anyhow` owo-colors = { version = "4.0.0" } # Colourised output compaction = { path = "../compaction" } chrono = { version = "0.4.26" } # Log helper tokio = { version = "1.20.1", features = ["full"] } # Async runtime -aws-config = { version = "1.2.0" } # Credential loading -aws-credential-types = { version = "1.2.0" } # Credential provider types -aws-types = { version = "1.2.0" } # for Region url = { version = "2.4.0" } diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index 14690903f1..8126a00165 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -14,8 +14,6 @@ * limitations under the License. */ -use aws_config::BehaviorVersion; -use aws_credential_types::provider::ProvideCredentials; use chrono::Local; use clap::Parser; use compaction::{merge_sorted_files, CompactionInput}; @@ -91,16 +89,6 @@ async fn main() -> color_eyre::Result<()> { let output_url = Url::parse(&args.output) .or_else(|_e| Url::parse(&("file://".to_owned() + &args.output)))?; - let config = aws_config::defaults(BehaviorVersion::latest()).load().await; - let region = config - .region() - .ok_or(color_eyre::eyre::eyre!("Can't determine AWS region"))?; - let creds: aws_credential_types::Credentials = config - .credentials_provider() - .ok_or(color_eyre::eyre::eyre!("Can't load AWS config"))? - .provide_credentials() - .await?; - let details = CompactionInput { input_files: input_urls, output_file: output_url, @@ -118,6 +106,6 @@ async fn main() -> color_eyre::Result<()> { sort_key_cols: vec![], }; - merge_sorted_files(Some(creds), region, &details).await?; + merge_sorted_files(&details).await?; Ok(()) } From 8e9595c57c19354b43f813a742f4e92bc9aaa5e0 Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 23 Apr 2024 16:18:58 +0000 Subject: [PATCH 047/129] WIP datafusion code --- rust/compaction/src/datafusion.rs | 19 ++++++++++++++++++- rust/compaction/src/lib.rs | 8 ++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 3e017f5876..2904ac1193 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -18,7 +18,10 @@ * limitations under the License. */ use crate::{aws_s3::ObjectStoreFactory, CompactionInput, CompactionResult}; -use datafusion::error::DataFusionError; +use datafusion::{ + error::DataFusionError, + execution::{config::SessionConfig, context::SessionContext}, +}; use url::Url; pub async fn compact( @@ -27,6 +30,20 @@ pub async fn compact( input_paths: &[Url], output_path: &Url, ) -> Result { + let mut sf = SessionConfig::new(); + sf.options_mut().execution.parquet.data_pagesize_limit = input_data.max_page_size; + sf.options_mut().execution.parquet.writer_version = input_data.writer_version.clone(); + sf.options_mut().execution.parquet.compression = Some(input_data.compression.clone()); + sf.options_mut().execution.parquet.max_statistics_size = Some(input_data.stats_truncate_length); + sf.options_mut().execution.parquet.max_row_group_size = input_data.max_row_group_size; + sf.options_mut() + .execution + .parquet + .column_index_truncate_length = Some(input_data.column_truncate_length); + + // Create my scalar UDF + let ctx = SessionContext::new_with_config(sf); + Ok(CompactionResult { rows_read: 0, rows_written: 0, diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 76fe58152e..b80177d07b 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -248,7 +248,7 @@ pub extern "C" fn ffi_merge_sorted_files( ) -> c_int { maybe_cfg_log(); if let Err(e) = color_eyre::install() { - warn!("Couldn't install color_eyre error handler"); + warn!("Couldn't install color_eyre error handler {e}"); } let Some(params) = (unsafe { input_data.as_ref() }) else { error!("input data pointer is null"); @@ -262,7 +262,7 @@ pub extern "C" fn ffi_merge_sorted_files( { Ok(v) => v, Err(e) => { - error!("Couldn't create Rust tokio runtime {}", e); + error!("Couldn't create Rust tokio runtime {e}"); return EIO; } }; @@ -270,7 +270,7 @@ pub extern "C" fn ffi_merge_sorted_files( let details = match TryInto::::try_into(params) { Ok(d) => d, Err(e) => { - error!("Couldn't convert compaction input data {}", e); + error!("Couldn't convert compaction input data {e}"); return EINVAL; } }; @@ -289,7 +289,7 @@ pub extern "C" fn ffi_merge_sorted_files( 0 } Err(e) => { - error!("merging error {}", e); + error!("merging error {e}"); -1 } } From d21eb19e0388764c6e4c046b44ec78980e041964 Mon Sep 17 00:00:00 2001 From: m09526 Date: Wed, 24 Apr 2024 13:59:02 +0000 Subject: [PATCH 048/129] DataFusion done --- rust/compaction/src/aws_s3.rs | 19 ++-- rust/compaction/src/datafusion.rs | 173 ++++++++++++++++++++++++++++-- rust/compaction/src/sketch.rs | 11 +- rust/compactor/src/bin/main.rs | 18 ++-- 4 files changed, 181 insertions(+), 40 deletions(-) diff --git a/rust/compaction/src/aws_s3.rs b/rust/compaction/src/aws_s3.rs index af12f7bc24..810348e2db 100644 --- a/rust/compaction/src/aws_s3.rs +++ b/rust/compaction/src/aws_s3.rs @@ -25,9 +25,9 @@ use std::{ sync::{Arc, Mutex}, }; -use arrow::error::ArrowError; use aws_types::region::Region; use bytes::Bytes; +use color_eyre::eyre::eyre; use futures::{stream::BoxStream, Future}; use log::info; use num_format::{Locale, ToFormattedString}; @@ -111,7 +111,7 @@ impl ObjectStoreFactory { /// # Errors /// /// If no credentials have been provided, then trying to access S3 URLs will fail. - pub fn get_object_store(&self, src: &Url) -> Result, ArrowError> { + pub fn get_object_store(&self, src: &Url) -> color_eyre::Result> { let scheme = src.scheme(); let mut borrow = self.store_map.borrow_mut(); // Perform a single lookup into the cache map @@ -135,29 +135,24 @@ impl ObjectStoreFactory { /// # Errors /// /// If no credentials have been provided, then trying to access S3 URLs will fail. - fn make_object_store(&self, src: &Url) -> Result, ArrowError> { + fn make_object_store(&self, src: &Url) -> color_eyre::Result> { match src.scheme() { "s3" => { if let Some(creds) = &self.creds { Ok(AmazonS3Builder::from_env() .with_credentials(creds.clone()) .with_region(self.region.as_ref()) - .with_bucket_name(src.host_str().ok_or( - ArrowError::InvalidArgumentError("invalid S3 bucket name".into()), - )?) + .with_bucket_name(src.host_str().ok_or(eyre!("invalid S3 bucket name"))?) .build() - .map(|e| Arc::new(LoggingObjectStore::new(Arc::new(e)))) - .map_err(|e| ArrowError::ExternalError(Box::new(e)))?) + .map(|e| Arc::new(LoggingObjectStore::new(Arc::new(e))))?) } else { - Err(ArrowError::InvalidArgumentError("Can't create AWS S3 object_store: no credentials provided to ObjectStoreFactory::from".into())) + Err(eyre!("Can't create AWS S3 object_store: no credentials provided to ObjectStoreFactory::from")) } } "file" => Ok(Arc::new(LoggingObjectStore::new(Arc::new( LocalFileSystem::new(), )))), - _ => Err(ArrowError::InvalidArgumentError( - "no object store for given schema".into(), - )), + _ => Err(eyre!("no object store for given schema")), } } } diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 2904ac1193..fcc8628a4b 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + /// `DataFusion` contains the implementation for performing Sleeper compactions /// using Apache `DataFusion`. /// @@ -17,35 +19,184 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -use crate::{aws_s3::ObjectStoreFactory, CompactionInput, CompactionResult}; +use crate::{ + aws_s3::{CountingObjectStore, ObjectStoreFactory}, + CompactionInput, CompactionResult, +}; +use arrow::{array::RecordBatch, util::pretty::pretty_format_batches}; use datafusion::{ + dataframe::DataFrameWriteOptions, error::DataFusionError, - execution::{config::SessionConfig, context::SessionContext}, + execution::{config::SessionConfig, context::SessionContext, options::ParquetReadOptions}, + parquet::basic::{BrotliLevel, GzipLevel, ZstdLevel}, + prelude::*, }; +use log::{error, info}; +use num_format::{Locale, ToFormattedString}; use url::Url; +/// Starts a Sleeper compaction. +/// +/// The object store factory must be able to produce an [`ObjectStore`] capable of reading +/// from the input URLs and writing to the output URL. A sketch file will be produced for +/// the output file. pub async fn compact( store_factory: &ObjectStoreFactory, input_data: &CompactionInput, input_paths: &[Url], output_path: &Url, ) -> Result { + info!("DataFusion compaction of {input_paths:?}"); + + let sf = create_session_cfg(input_data); + let ctx = SessionContext::new_with_config(sf); + + // Register some objects store from first object + let store = register_store(store_factory, input_paths, output_path, &ctx)?; + + // Sort on row key columns then sort columns (nulls last) + let sort_order = sort_order(input_data); + info!("Row key and sort column order {sort_order:?}"); + + let po = ParquetReadOptions::default().file_sort_order(vec![sort_order.clone()]); + let frame = ctx.read_parquet(input_paths.to_owned(), po).await?; + + // Extract all column names + let col_names = frame.schema().clone().strip_qualifiers().field_names(); + info!("All columns in schema {col_names:?}"); + + // Perform projection of all columns + let col_names_expr = frame + .schema() + .clone() + .strip_qualifiers() + .field_names() + .iter() + .map(|x| col(x)) + .collect::>(); + let frame = frame.sort(sort_order)?.select(col_names_expr)?; + + // Show explanation of plan + let explained = frame.clone().explain(false, false)?.collect().await?; + let output = pretty_format_batches(&explained)?; + info!("DataFusion plan:\n {output}"); + + let mut pqo = ctx.copied_table_options().parquet; + for col in &col_names { + let col_opts = pqo.column_specific_options.entry(col.into()).or_default(); + let dict_encode = if input_data.dict_enc_row_keys && input_data.row_key_cols.contains(col) { + println!("{col} is a row key and row key dictionary enabled"); + true + } else if input_data.dict_enc_sort_keys && input_data.sort_key_cols.contains(col) { + println!("{col} is a sort column and sort column dictionary encoding enabled"); + true + } else if input_data.dict_enc_values { + println!("{col} is a value and value dictionary encoding enabled"); + true + } else { + false + }; + col_opts.dictionary_enabled = Some(dict_encode); + } + + let result = frame + .write_parquet( + output_path.as_str(), + DataFrameWriteOptions::new(), + Some(pqo), + ) + .await?; + + info!( + "Object store read {} bytes from {} GETs", + store + .get_bytes_read() + .unwrap_or(0) + .to_formatted_string(&Locale::en), + store + .get_count() + .unwrap_or(0) + .to_formatted_string(&Locale::en) + ); + let rows_written = result.iter().map(RecordBatch::num_rows).sum::(); + // The rows read will be same as rows_written. + Ok(CompactionResult { + rows_read: rows_written, + rows_written, + }) + // TODO: Sketches +} + +/// Convert a Sleeper compression codec string to one `DataFusion` understands. +fn get_compression(compression: &str) -> String { + match compression.to_lowercase().as_str() { + x @ "uncompressed" | x @ "snappy" | x @ "lzo" | x @ "lz4" => x.into(), + "gzip" => format!("gzip({})", GzipLevel::default().compression_level()), + "brotli" => format!("brotli({})", BrotliLevel::default().compression_level()), + "zstd" => format!("zstd({})", ZstdLevel::default().compression_level()), + _ => { + error!("Unknown compression"); + unimplemented!() + } + } +} + +/// Create the `DataFusion` session configuration for a given compaction. +/// +/// This sets as many parameters as possible from the given input data. +/// +fn create_session_cfg(input_data: &CompactionInput) -> SessionConfig { let mut sf = SessionConfig::new(); + sf.options_mut().execution.parquet.max_row_group_size = input_data.max_row_group_size; sf.options_mut().execution.parquet.data_pagesize_limit = input_data.max_page_size; + sf.options_mut().execution.parquet.compression = Some(get_compression(&input_data.compression)); sf.options_mut().execution.parquet.writer_version = input_data.writer_version.clone(); - sf.options_mut().execution.parquet.compression = Some(input_data.compression.clone()); - sf.options_mut().execution.parquet.max_statistics_size = Some(input_data.stats_truncate_length); - sf.options_mut().execution.parquet.max_row_group_size = input_data.max_row_group_size; sf.options_mut() .execution .parquet .column_index_truncate_length = Some(input_data.column_truncate_length); + sf.options_mut().execution.parquet.max_statistics_size = Some(input_data.stats_truncate_length); + sf +} - // Create my scalar UDF - let ctx = SessionContext::new_with_config(sf); +/// Creates the sort order for a given schema. +/// +/// This is a list of the row key columns followed by the sort columns. +/// +fn sort_order(input_data: &CompactionInput) -> Vec { + let sort_order = input_data + .row_key_cols + .iter() + .chain(input_data.sort_key_cols.iter()) + .map(|s| col(s).sort(true, false)) + .collect::>(); + sort_order +} - Ok(CompactionResult { - rows_read: 0, - rows_written: 0, - }) +/// Takes the first Url in the input_paths list and the output path +/// and registers the appropriate [`ObjectStore`] for it. +/// +/// DataFusion doesn't seem to like loading a single file set from different object stores +/// so we only register the first one. +/// +/// # Errors +/// If we can't create an [`ObjectStore`] for a known URL then this will fail. +/// +fn register_store( + store_factory: &ObjectStoreFactory, + input_paths: &[Url], + output_path: &Url, + ctx: &SessionContext, +) -> Result, DataFusionError> { + let store = store_factory + .get_object_store(&input_paths[0]) + .map_err(|e| DataFusionError::External(e.into()))?; + ctx.runtime_env() + .register_object_store(&input_paths[0], store.clone().as_object_store()); + let store = store_factory + .get_object_store(output_path) + .map_err(|e| DataFusionError::External(e.into()))?; + ctx.runtime_env() + .register_object_store(output_path, store.clone().as_object_store()); + Ok(store) } diff --git a/rust/compaction/src/sketch.rs b/rust/compaction/src/sketch.rs index 9bb6512b70..cb3c3fd6fb 100644 --- a/rust/compaction/src/sketch.rs +++ b/rust/compaction/src/sketch.rs @@ -20,7 +20,6 @@ use arrow::array::{ArrayAccessor, AsArray}; use arrow::datatypes::{ BinaryType, DataType, Int32Type, Int64Type, LargeBinaryType, LargeUtf8Type, Schema, Utf8Type, }; -use arrow::error::ArrowError; use arrow::record_batch::RecordBatch; use bytes::{BufMut, Bytes}; use cxx::{Exception, UniquePtr}; @@ -267,7 +266,6 @@ impl DataSketchVariant { /// followed by the bytes of the sketch. /// /// # Errors -/// Any I/O errors are wrapped into an [`ArrowError`] if thrown. /// The data sketch serialisation might also throw errors from the underlying /// data sketch library. #[allow(clippy::cast_possible_truncation)] @@ -275,15 +273,13 @@ pub fn serialise_sketches( store_factory: &ObjectStoreFactory, path: &Url, sketches: &[DataSketchVariant], -) -> Result<(), ArrowError> { +) -> color_eyre::Result<()> { let mut buf = vec![].writer(); let mut size = 0; // for each sketch write the size i32, followed by bytes for sketch in sketches { - let serialised = sketch - .serialize(0) - .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + let serialised = sketch.serialize(0)?; buf.write_all(&(serialised.len() as u32).to_be_bytes())?; buf.write_all(&serialised)?; size += serialised.len() + size_of::(); @@ -295,8 +291,7 @@ pub fn serialise_sketches( // Save to object store let store = store_factory.get_object_store(path)?; - futures::executor::block_on(store.put(&store_path, Bytes::from(buf.into_inner()))) - .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + futures::executor::block_on(store.put(&store_path, Bytes::from(buf.into_inner())))?; info!( "Serialised {} ({} bytes) sketches to {}", diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index 8126a00165..2ea7b0d862 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -42,15 +42,15 @@ struct CmdLineArgs { /// List of input Parquet files (must be sorted) as URLs #[arg(num_args=1.., required=true)] input: Vec, - // Column number for a row key fields - #[arg(short = 'k', long, default_value = "0")] - row_keys: Vec, - // Column number for sort columns - #[arg(short = 's', long, default_value = "0")] - sort_column: Vec, + /// Column names for a row key fields + #[arg(short = 'k', long, num_args=1.., required=true)] + row_keys: Vec, + /// Column names for sort columns + #[arg(short = 's', long)] + sort_column: Vec, } -#[tokio::main] +#[tokio::main(flavor = "multi_thread")] async fn main() -> color_eyre::Result<()> { // Install coloured errors color_eyre::install().unwrap(); @@ -102,8 +102,8 @@ async fn main() -> color_eyre::Result<()> { dict_enc_sort_keys: true, dict_enc_values: true, region: HashMap::default(), - row_key_cols: vec!["key".into()], - sort_key_cols: vec![], + row_key_cols: args.row_keys, + sort_key_cols: args.sort_column, }; merge_sorted_files(&details).await?; From f4ac9bec173e3fc18e3d5ba5ff6118a2bedf8da0 Mon Sep 17 00:00:00 2001 From: m09526 Date: Wed, 24 Apr 2024 14:09:30 +0000 Subject: [PATCH 049/129] clippy warnings --- rust/compaction/src/datafusion.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index fcc8628a4b..89225642aa 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -72,7 +72,7 @@ pub async fn compact( .strip_qualifiers() .field_names() .iter() - .map(|x| col(x)) + .map(col) .collect::>(); let frame = frame.sort(sort_order)?.select(col_names_expr)?; @@ -130,7 +130,7 @@ pub async fn compact( /// Convert a Sleeper compression codec string to one `DataFusion` understands. fn get_compression(compression: &str) -> String { match compression.to_lowercase().as_str() { - x @ "uncompressed" | x @ "snappy" | x @ "lzo" | x @ "lz4" => x.into(), + x @ ("uncompressed" | "snappy" | "lzo" | "lz4") => x.into(), "gzip" => format!("gzip({})", GzipLevel::default().compression_level()), "brotli" => format!("brotli({})", BrotliLevel::default().compression_level()), "zstd" => format!("zstd({})", ZstdLevel::default().compression_level()), @@ -173,10 +173,10 @@ fn sort_order(input_data: &CompactionInput) -> Vec { sort_order } -/// Takes the first Url in the input_paths list and the output path +/// Takes the first Url in `input_paths` list and `output_path` /// and registers the appropriate [`ObjectStore`] for it. /// -/// DataFusion doesn't seem to like loading a single file set from different object stores +/// `DataFusion` doesn't seem to like loading a single file set from different object stores /// so we only register the first one. /// /// # Errors From 4cf2ffcda5ba4ad619fba8fd6788642f18813c92 Mon Sep 17 00:00:00 2001 From: m09526 Date: Wed, 24 Apr 2024 15:20:34 +0000 Subject: [PATCH 050/129] Compaction links up --- .../compaction/jobexecution/RustBridge.java | 26 ++++++++++------- rust/Cargo.lock | 1 + rust/Cargo.toml | 2 +- rust/compaction/src/datafusion.rs | 29 ++++++++----------- rust/compactor/Cargo.toml | 1 + rust/compactor/src/bin/main.rs | 11 ++++++- 6 files changed, 41 insertions(+), 29 deletions(-) diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java index dac1ee74a3..971e47d33f 100644 --- a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java @@ -244,19 +244,25 @@ public void populate(final T[] arr) { final jnr.ffi.Runtime r = len.struct().getRuntime(); // Calculate size needed for array of pointers int ptrSize = r.findType(NativeType.ADDRESS).size(); - int size = arr.length * ptrSize; - // Allocate some memory for string pointers - this.basePtr = r.getMemoryManager().allocateDirect(size); - this.arrayBase.set(basePtr); + // Null out zero length arrays + if (arr.length > 0) { + int size = arr.length * ptrSize; + // Allocate some memory for string pointers + this.basePtr = r.getMemoryManager().allocateDirect(size); + this.arrayBase.set(basePtr); - this.items = new jnr.ffi.Pointer[arr.length]; + this.items = new jnr.ffi.Pointer[arr.length]; - for (int i = 0; i < arr.length; i++) { - setValue(arr[i], i, r); - } + for (int i = 0; i < arr.length; i++) { + setValue(arr[i], i, r); + } - // Bulk set the pointers in the base array - this.basePtr.put(0, this.items, 0, this.items.length); + // Bulk set the pointers in the base array + this.basePtr.put(0, this.items, 0, this.items.length); + } else { + this.basePtr = null; + this.items = null; + } // Set length of array in struct this.len.set(arr.length); diff --git a/rust/Cargo.lock b/rust/Cargo.lock index af0f6504be..ccffd28410 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1050,6 +1050,7 @@ dependencies = [ "env_logger", "human-panic", "log", + "num-format", "owo-colors 4.0.0", "thiserror", "tokio", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 1b0a27d719..dc086aa607 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -19,4 +19,4 @@ resolver = "2" incremental = true lto = false overflow-checks = true -debug = true +debug = false diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 89225642aa..ea6003e7d4 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -82,20 +82,15 @@ pub async fn compact( info!("DataFusion plan:\n {output}"); let mut pqo = ctx.copied_table_options().parquet; + // Figure out which columns should be dictionary encoded for col in &col_names { let col_opts = pqo.column_specific_options.entry(col.into()).or_default(); - let dict_encode = if input_data.dict_enc_row_keys && input_data.row_key_cols.contains(col) { - println!("{col} is a row key and row key dictionary enabled"); - true - } else if input_data.dict_enc_sort_keys && input_data.sort_key_cols.contains(col) { - println!("{col} is a sort column and sort column dictionary encoding enabled"); - true - } else if input_data.dict_enc_values { - println!("{col} is a value and value dictionary encoding enabled"); - true - } else { - false - }; + let dict_encode = (input_data.dict_enc_row_keys && input_data.row_key_cols.contains(col)) + || (input_data.dict_enc_sort_keys && input_data.sort_key_cols.contains(col)) + // Check value columns + || (input_data.dict_enc_values + && !input_data.row_key_cols.contains(col) + && !input_data.sort_key_cols.contains(col)); col_opts.dictionary_enabled = Some(dict_encode); } @@ -188,15 +183,15 @@ fn register_store( output_path: &Url, ctx: &SessionContext, ) -> Result, DataFusionError> { - let store = store_factory + let in_store = store_factory .get_object_store(&input_paths[0]) .map_err(|e| DataFusionError::External(e.into()))?; ctx.runtime_env() - .register_object_store(&input_paths[0], store.clone().as_object_store()); - let store = store_factory + .register_object_store(&input_paths[0], in_store.clone().as_object_store()); + let out_store = store_factory .get_object_store(output_path) .map_err(|e| DataFusionError::External(e.into()))?; ctx.runtime_env() - .register_object_store(output_path, store.clone().as_object_store()); - Ok(store) + .register_object_store(output_path, out_store.clone().as_object_store()); + Ok(in_store) } diff --git a/rust/compactor/Cargo.toml b/rust/compactor/Cargo.toml index 447792e603..085dc82347 100644 --- a/rust/compactor/Cargo.toml +++ b/rust/compactor/Cargo.toml @@ -40,3 +40,4 @@ compaction = { path = "../compaction" } chrono = { version = "0.4.26" } # Log helper tokio = { version = "1.20.1", features = ["full"] } # Async runtime url = { version = "2.4.0" } +num-format = { version = "0.4.4" } # Formatted numbers diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index 2ea7b0d862..90cd0dc04e 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -18,6 +18,8 @@ use chrono::Local; use clap::Parser; use compaction::{merge_sorted_files, CompactionInput}; use human_panic::setup_panic; +use log::info; +use num_format::{Locale, ToFormattedString}; use std::{collections::HashMap, io::Write}; use url::Url; @@ -106,6 +108,13 @@ async fn main() -> color_eyre::Result<()> { sort_key_cols: args.sort_column, }; - merge_sorted_files(&details).await?; + let result = merge_sorted_files(&details).await; + if let Ok(data) = result { + info!( + "Compaction read {} rows and wrote {} rows", + data.rows_read.to_formatted_string(&Locale::en), + data.rows_written.to_formatted_string(&Locale::en) + ); + } Ok(()) } From e55e8f7f00c57aa6f40b0a1de7af27ff6116e888 Mon Sep 17 00:00:00 2001 From: m09526 Date: Wed, 24 Apr 2024 16:18:22 +0000 Subject: [PATCH 051/129] Partitions work --- rust/compaction/src/datafusion.rs | 61 +++++++++++++++++++++++++++---- rust/compaction/src/details.rs | 8 ++-- rust/compaction/src/lib.rs | 20 ++++------ rust/compactor/src/bin/main.rs | 14 ++++++- 4 files changed, 77 insertions(+), 26 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index ea6003e7d4..c86993ecbd 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -1,5 +1,3 @@ -use std::sync::Arc; - /// `DataFusion` contains the implementation for performing Sleeper compactions /// using Apache `DataFusion`. /// @@ -21,7 +19,7 @@ use std::sync::Arc; */ use crate::{ aws_s3::{CountingObjectStore, ObjectStoreFactory}, - CompactionInput, CompactionResult, + ColRange, CompactionInput, CompactionResult, PartitionBound, }; use arrow::{array::RecordBatch, util::pretty::pretty_format_batches}; use datafusion::{ @@ -33,6 +31,7 @@ use datafusion::{ }; use log::{error, info}; use num_format::{Locale, ToFormattedString}; +use std::{collections::HashMap, sync::Arc}; use url::Url; /// Starts a Sleeper compaction. @@ -47,7 +46,7 @@ pub async fn compact( output_path: &Url, ) -> Result { info!("DataFusion compaction of {input_paths:?}"); - + info!("Compaction region {:?}", input_data.region); let sf = create_session_cfg(input_data); let ctx = SessionContext::new_with_config(sf); @@ -59,7 +58,7 @@ pub async fn compact( info!("Row key and sort column order {sort_order:?}"); let po = ParquetReadOptions::default().file_sort_order(vec![sort_order.clone()]); - let frame = ctx.read_parquet(input_paths.to_owned(), po).await?; + let mut frame = ctx.read_parquet(input_paths.to_owned(), po).await?; // Extract all column names let col_names = frame.schema().clone().strip_qualifiers().field_names(); @@ -74,7 +73,12 @@ pub async fn compact( .iter() .map(col) .collect::>(); - let frame = frame.sort(sort_order)?.select(col_names_expr)?; + + // Create plan + if let Some(expr) = region_filter(&input_data.region) { + frame = frame.filter(expr)?; + } + frame = frame.sort(sort_order)?.select(col_names_expr)?; // Show explanation of plan let explained = frame.clone().explain(false, false)?.collect().await?; @@ -119,7 +123,50 @@ pub async fn compact( rows_read: rows_written, rows_written, }) - // TODO: Sketches +} + +fn region_filter(region: &HashMap) -> Option { + let mut col_exprs = vec![]; + for (name, range) in region { + let min_bound = match range.lower { + PartitionBound::Int32(val) => lit(val), + PartitionBound::Int64(val) => lit(val), + PartitionBound::String(val) => lit(val.to_owned()), + PartitionBound::ByteArray(val) => lit(val.to_owned()), + }; + let lower_expr = if range.lower_inclusive { + col(name).gt_eq(min_bound) + } else { + col(name).gt(min_bound) + }; + + let max_bound = match range.upper { + PartitionBound::Int32(val) => lit(val), + PartitionBound::Int64(val) => lit(val), + PartitionBound::String(val) => lit(val.to_owned()), + PartitionBound::ByteArray(val) => lit(val.to_owned()), + }; + let upper_expr = if range.upper_inclusive { + col(name).lt_eq(max_bound) + } else { + col(name).lt(max_bound) + }; + + let expr = lower_expr.and(upper_expr); + col_exprs.push(expr); + } + + // join them together + // TODO: write this more Rust like + if !col_exprs.is_empty() { + let mut expr = col_exprs[0].clone(); + for idx in 1..col_exprs.len() { + expr = expr.and(col_exprs[idx].clone()); + } + Some(expr) + } else { + None + } } /// Convert a Sleeper compression codec string to one `DataFusion` understands. diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index be85b8d4d9..f961786daf 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -26,10 +26,10 @@ use url::Url; /// Type safe variant for Sleeper partition boundary #[derive(Debug, Copy, Clone)] pub enum PartitionBound { - Int32 { val: i32 }, - Int64 { val: i64 }, - String { val: &'static str }, - ByteArray { val: &'static [i8] }, + Int32(i32), + Int64(i64), + String(&'static str), + ByteArray(&'static [u8]), } /// All the information for a a Sleeper compaction. diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index b80177d07b..2262fd5df7 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -375,12 +375,8 @@ fn unpack_variant_array( }) .zip(schema_types.iter()) .map(|(&bptr, type_id)| match type_id { - 1 => Ok(PartitionBound::Int32 { - val: unsafe { *bptr.cast::() }, - }), - 2 => Ok(PartitionBound::Int64 { - val: unsafe { *bptr.cast::() }, - }), + 1 => Ok(PartitionBound::Int32(unsafe { *bptr.cast::() })), + 2 => Ok(PartitionBound::Int64(unsafe { *bptr.cast::() })), 3 => { //unpack length (signed because it's from Java) let str_len = unsafe { *bptr.cast::() }; @@ -392,7 +388,7 @@ fn unpack_variant_array( #[allow(clippy::cast_sign_loss)] slice::from_raw_parts(bptr.byte_add(4).cast::(), str_len as usize) }) - .map(|v| PartitionBound::String { val: v }) + .map(|v| PartitionBound::String(v)) } 4 => { //unpack length (signed because it's from Java) @@ -401,12 +397,10 @@ fn unpack_variant_array( error!("Illegal byte array length in FFI array: {byte_len}"); panic!("Illegal byte array length in FFI array: {byte_len}"); } - Ok(PartitionBound::ByteArray { - val: unsafe { - #[allow(clippy::cast_sign_loss)] - slice::from_raw_parts(bptr.byte_add(4).cast::(), byte_len as usize) - }, - }) + Ok(PartitionBound::ByteArray(unsafe { + #[allow(clippy::cast_sign_loss)] + slice::from_raw_parts(bptr.byte_add(4).cast::(), byte_len as usize) + })) } x => { error!("Unexpected type id {x}"); diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index 90cd0dc04e..b51ee267e4 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -16,7 +16,7 @@ use chrono::Local; use clap::Parser; -use compaction::{merge_sorted_files, CompactionInput}; +use compaction::{merge_sorted_files, ColRange, CompactionInput}; use human_panic::setup_panic; use log::info; use num_format::{Locale, ToFormattedString}; @@ -91,6 +91,16 @@ async fn main() -> color_eyre::Result<()> { let output_url = Url::parse(&args.output) .or_else(|_e| Url::parse(&("file://".to_owned() + &args.output)))?; + let mut map = HashMap::new(); + map.insert( + "key".into(), + ColRange { + lower: compaction::PartitionBound::String("h"), + lower_inclusive: true, + upper: compaction::PartitionBound::String("m"), + upper_inclusive: false, + }, + ); let details = CompactionInput { input_files: input_urls, output_file: output_url, @@ -103,7 +113,7 @@ async fn main() -> color_eyre::Result<()> { dict_enc_row_keys: true, dict_enc_sort_keys: true, dict_enc_values: true, - region: HashMap::default(), + region: map, row_key_cols: args.row_keys, sort_key_cols: args.sort_column, }; From 72922e8a17e94734983e4c2d294dd0b40c2f938c Mon Sep 17 00:00:00 2001 From: m09526 Date: Thu, 25 Apr 2024 16:16:22 +0000 Subject: [PATCH 052/129] better region coding --- rust/compaction/src/datafusion.rs | 148 ++++++++++++++++++++---------- 1 file changed, 99 insertions(+), 49 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index c86993ecbd..7f94ef16de 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -21,12 +21,15 @@ use crate::{ aws_s3::{CountingObjectStore, ObjectStoreFactory}, ColRange, CompactionInput, CompactionResult, PartitionBound, }; -use arrow::{array::RecordBatch, util::pretty::pretty_format_batches}; +use arrow::{array::RecordBatch, error::ArrowError, util::pretty::pretty_format_batches}; use datafusion::{ dataframe::DataFrameWriteOptions, error::DataFusionError, execution::{config::SessionConfig, context::SessionContext, options::ParquetReadOptions}, - parquet::basic::{BrotliLevel, GzipLevel, ZstdLevel}, + parquet::{ + arrow::{async_reader::ParquetObjectReader, ParquetRecordBatchStreamBuilder}, + basic::{BrotliLevel, GzipLevel, ZstdLevel}, + }, prelude::*, }; use log::{error, info}; @@ -50,13 +53,14 @@ pub async fn compact( let sf = create_session_cfg(input_data); let ctx = SessionContext::new_with_config(sf); - // Register some objects store from first object + // Register some object store from first input file and output file let store = register_store(store_factory, input_paths, output_path, &ctx)?; // Sort on row key columns then sort columns (nulls last) let sort_order = sort_order(input_data); info!("Row key and sort column order {sort_order:?}"); + // Tell DataFusion that the row key columns and sort columns are already sorted let po = ParquetReadOptions::default().file_sort_order(vec![sort_order.clone()]); let mut frame = ctx.read_parquet(input_paths.to_owned(), po).await?; @@ -64,7 +68,12 @@ pub async fn compact( let col_names = frame.schema().clone().strip_qualifiers().field_names(); info!("All columns in schema {col_names:?}"); - // Perform projection of all columns + // If we have a partition region, apply it first + if let Some(expr) = region_filter(&input_data.region) { + frame = frame.filter(expr)?; + } + + // Perform sort of row key and sort columns and projection of all columns let col_names_expr = frame .schema() .clone() @@ -73,11 +82,6 @@ pub async fn compact( .iter() .map(col) .collect::>(); - - // Create plan - if let Some(expr) = region_filter(&input_data.region) { - frame = frame.filter(expr)?; - } frame = frame.sort(sort_order)?.select(col_names_expr)?; // Show explanation of plan @@ -98,7 +102,7 @@ pub async fn compact( col_opts.dictionary_enabled = Some(dict_encode); } - let result = frame + let _ = frame .write_parquet( output_path.as_str(), DataFrameWriteOptions::new(), @@ -117,56 +121,102 @@ pub async fn compact( .unwrap_or(0) .to_formatted_string(&Locale::en) ); - let rows_written = result.iter().map(RecordBatch::num_rows).sum::(); + + // Find rows just written to output Parquet file + let pq_reader = get_parquet_builder(store_factory, output_path) + .await + .map_err(|e| DataFusionError::External(e.into()))?; + let num_rows = pq_reader.metadata().file_metadata().num_rows(); + // The rows read will be same as rows_written. Ok(CompactionResult { - rows_read: rows_written, - rows_written, + rows_read: num_rows as usize, + rows_written: num_rows as usize, }) } +/// Create an asynchronous builder for reading Parquet files from an object store. +/// +/// The URL must start with a scheme that the object store recognises, e.g. "file" or "s3". +/// +/// # Errors +/// This function will return an error if it couldn't connect to S3 or open a valid +/// Parquet file. +pub async fn get_parquet_builder( + store_factory: &ObjectStoreFactory, + src: &Url, +) -> color_eyre::Result> { + let store = store_factory.get_object_store(src)?; + // HEAD the file to get metadata + let path = object_store::path::Path::from(src.path()); + let object_meta = store + .head(&path) + .await + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + // Create a reader for the target file, use it to construct a Stream + let reader = ParquetObjectReader::new(store.as_object_store(), object_meta); + Ok(ParquetRecordBatchStreamBuilder::new(reader).await?) +} + +/// Create the `DataFusion` filtering expression from a Sleeper region. +/// +/// For each column in the row keys, we look up the partition range for that +/// column and create a expression tree that combines all the various filtering conditions. +/// fn region_filter(region: &HashMap) -> Option { - let mut col_exprs = vec![]; + let mut col_expr: Option = None; for (name, range) in region { - let min_bound = match range.lower { - PartitionBound::Int32(val) => lit(val), - PartitionBound::Int64(val) => lit(val), - PartitionBound::String(val) => lit(val.to_owned()), - PartitionBound::ByteArray(val) => lit(val.to_owned()), - }; - let lower_expr = if range.lower_inclusive { - col(name).gt_eq(min_bound) - } else { - col(name).gt(min_bound) - }; - - let max_bound = match range.upper { - PartitionBound::Int32(val) => lit(val), - PartitionBound::Int64(val) => lit(val), - PartitionBound::String(val) => lit(val.to_owned()), - PartitionBound::ByteArray(val) => lit(val.to_owned()), - }; - let upper_expr = if range.upper_inclusive { - col(name).lt_eq(max_bound) - } else { - col(name).lt(max_bound) - }; - + let lower_expr = lower_bound_expr(range, name); + let upper_expr = upper_bound_expr(range, name); let expr = lower_expr.and(upper_expr); - col_exprs.push(expr); + // Combine this column filter with any previous column filter + col_expr = match col_expr { + Some(original) => Some(original.and(expr)), + None => Some(expr), + } } + col_expr +} - // join them together - // TODO: write this more Rust like - if !col_exprs.is_empty() { - let mut expr = col_exprs[0].clone(); - for idx in 1..col_exprs.len() { - expr = expr.and(col_exprs[idx].clone()); - } - Some(expr) +/// Calculate the upper bound expression on a given [`ColRange`]. +/// +/// This takes into account the inclusive/exclusive nature of the bound. +/// +fn upper_bound_expr(range: &ColRange, name: &String) -> Expr { + let max_bound = bound_to_lit_expr(&range.upper); + let upper_expr = if range.upper_inclusive { + col(name).lt_eq(max_bound) } else { - None - } + col(name).lt(max_bound) + }; + upper_expr +} + +/// Calculate the lower bound expression on a given [`ColRange`]. +/// +/// This takes into account the inclusive/exclusive nature of the bound. +/// +fn lower_bound_expr(range: &ColRange, name: &String) -> Expr { + let min_bound = bound_to_lit_expr(&range.lower); + let lower_expr = if range.lower_inclusive { + col(name).gt_eq(min_bound) + } else { + col(name).gt(min_bound) + }; + lower_expr +} + +/// Convert a [`PartitionBound`] to an [`Expr`] that can be +/// used in a bigger expression. +/// +fn bound_to_lit_expr(bound: &PartitionBound) -> Expr { + let expr = match bound { + PartitionBound::Int32(val) => lit(*val), + PartitionBound::Int64(val) => lit(*val), + PartitionBound::String(val) => lit(val.to_owned()), + PartitionBound::ByteArray(val) => lit(val.to_owned()), + }; + expr } /// Convert a Sleeper compression codec string to one `DataFusion` understands. From 43a40f6c0a1282f713b2c83607737de5c09a6ab9 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 26 Apr 2024 11:22:12 +0000 Subject: [PATCH 053/129] Tidied some comments and test behaviour --- rust/Cargo.lock | 56 +++++++++--------------------- rust/compaction/src/datafusion.rs | 4 +-- rust/compaction/src/details.rs | 57 ++++++++++++++++++++++++------- rust/compaction/src/lib.rs | 14 ++++---- rust/compactor/src/bin/main.rs | 38 ++++++++++++++------- 5 files changed, 97 insertions(+), 72 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index ccffd28410..e070849a16 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -356,9 +356,9 @@ dependencies = [ [[package]] name = "async-compression" -version = "0.4.8" +version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07dbbf24db18d609b1462965249abdf49129ccad073ec257da372adc83259c60" +checksum = "4e9eabd7a98fe442131a17c316bd9349c43695e49e730c3c8e12cfb5f4da2693" dependencies = [ "bzip2", "flate2", @@ -1589,9 +1589,9 @@ dependencies = [ [[package]] name = "flate2" -version = "1.0.28" +version = "1.0.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46303f565772937ffe1d394a4fac6f411c6013172fadde9dcdb1e147a086940e" +checksum = "4556222738635b7a3417ae6130d8f52201e45a0c4d1a907f0826383adb5f85e7" dependencies = [ "crc32fast", "miniz_oxide", @@ -2212,9 +2212,9 @@ checksum = "01cda141df6706de531b6c46c3a33ecca755538219bd484262fa09410c13539c" [[package]] name = "lock_api" -version = "0.4.11" +version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" +checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" dependencies = [ "autocfg", "scopeguard", @@ -2529,9 +2529,9 @@ checksum = "caff54706df99d2a78a5a4e3455ff45448d81ef1bb63c22cd14052ca0e993a3f" [[package]] name = "parking_lot" -version = "0.12.1" +version = "0.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +checksum = "7e4af0ca4f6caed20e900d564c242b8e5d4903fdacf31d3daf527b66fe6f42fb" dependencies = [ "lock_api", "parking_lot_core", @@ -2539,15 +2539,15 @@ dependencies = [ [[package]] name = "parking_lot_core" -version = "0.9.9" +version = "0.9.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" +checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" dependencies = [ "cfg-if", "libc", "redox_syscall", "smallvec", - "windows-targets 0.48.5", + "windows-targets 0.52.5", ] [[package]] @@ -2744,11 +2744,11 @@ dependencies = [ [[package]] name = "redox_syscall" -version = "0.4.1" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +checksum = "469052894dcb553421e483e4209ee581a45100d31b4018de03e5a7ad86374a7e" dependencies = [ - "bitflags 1.3.2", + "bitflags 2.5.0", ] [[package]] @@ -3786,37 +3786,15 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "winapi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - [[package]] name = "winapi-util" -version = "0.1.6" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596" +checksum = "4d4cc384e1e73b93bafa6fb4f1df8c41695c8a91cf9c4c64358067d15a7b6c6b" dependencies = [ - "winapi", + "windows-sys 0.52.0", ] -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - [[package]] name = "windows-core" version = "0.52.0" diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 7f94ef16de..14b984ec84 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -21,7 +21,7 @@ use crate::{ aws_s3::{CountingObjectStore, ObjectStoreFactory}, ColRange, CompactionInput, CompactionResult, PartitionBound, }; -use arrow::{array::RecordBatch, error::ArrowError, util::pretty::pretty_format_batches}; +use arrow::{error::ArrowError, util::pretty::pretty_format_batches}; use datafusion::{ dataframe::DataFrameWriteOptions, error::DataFusionError, @@ -44,7 +44,7 @@ use url::Url; /// the output file. pub async fn compact( store_factory: &ObjectStoreFactory, - input_data: &CompactionInput, + input_data: &CompactionInput<'_>, input_paths: &[Url], output_path: &Url, ) -> Result { diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index f961786daf..c37a218dd5 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -25,16 +25,16 @@ use url::Url; /// Type safe variant for Sleeper partition boundary #[derive(Debug, Copy, Clone)] -pub enum PartitionBound { +pub enum PartitionBound<'a> { Int32(i32), Int64(i64), - String(&'static str), - ByteArray(&'static [u8]), + String(&'a str), + ByteArray(&'a [u8]), } /// All the information for a a Sleeper compaction. #[derive(Debug)] -pub struct CompactionInput { +pub struct CompactionInput<'a> { pub input_files: Vec, pub output_file: Url, pub row_key_cols: Vec, @@ -48,15 +48,36 @@ pub struct CompactionInput { pub dict_enc_row_keys: bool, pub dict_enc_sort_keys: bool, pub dict_enc_values: bool, - pub region: HashMap, + pub region: HashMap>, +} + +impl Default for CompactionInput<'_> { + fn default() -> Self { + Self { + input_files: Default::default(), + output_file: Url::parse("file:///").unwrap(), + row_key_cols: Default::default(), + sort_key_cols: Default::default(), + max_row_group_size: 1_000_000, + max_page_size: 65535, + compression: "zstd".into(), + writer_version: "2.0".into(), + column_truncate_length: usize::MAX, + stats_truncate_length: usize::MAX, + dict_enc_row_keys: true, + dict_enc_sort_keys: true, + dict_enc_values: true, + region: Default::default(), + } + } } /// Defines a partition range of a single column. #[derive(Debug, Copy, Clone)] -pub struct ColRange { - pub lower: PartitionBound, +pub struct ColRange<'a> { + pub lower: PartitionBound<'a>, pub lower_inclusive: bool, - pub upper: PartitionBound, + pub upper: PartitionBound<'a>, pub upper_inclusive: bool, } @@ -81,19 +102,31 @@ pub struct CompactionResult { /// /// # Examples /// ```no_run -/// # use arrow::error::ArrowError; /// # use url::Url; /// # use aws_types::region::Region; -/// # use crate::compaction::merge_sorted_files; +/// # use std::collections::HashMap; +/// # use crate::compaction::{merge_sorted_files, CompactionInput, PartitionBound, ColRange}; +/// let mut compaction_input = CompactionInput::default(); +/// compaction_input.input_files = vec![Url::parse("file:///path/to/file1.parquet").unwrap()]; +/// compaction_input.output_file = Url::parse("file:///path/to/output").unwrap(); +/// compaction_input.row_key_cols = vec!["key".into()]; +/// let mut region : HashMap> = HashMap::new(); +/// region.insert("key".into(), ColRange { +/// lower : PartitionBound::String("a"), +/// lower_inclusive: true, +/// upper: PartitionBound::String("h"), +/// upper_inclusive: true, +/// }); +/// /// # tokio_test::block_on(async { -/// let result = merge_sorted_files(None, &Region::new("eu-west-2"), &vec![Url::parse("file:///path/to/file1.parquet").unwrap()], &Url::parse("file:///path/to/output").unwrap(), 65535, 1_000_000, vec![0], vec![0]).await; +/// let result = merge_sorted_files(&compaction_input).await; /// # }) /// ``` /// /// # Errors /// There must be at least one input file. /// -pub async fn merge_sorted_files(input_data: &CompactionInput) -> Result { +pub async fn merge_sorted_files(input_data: &CompactionInput<'_>) -> Result { // Read the schema from the first file if input_data.input_files.is_empty() { Err(eyre!("No input paths supplied")) diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 2262fd5df7..4cf1de34bb 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -108,10 +108,10 @@ pub struct FFICompactionParams { region_maxs_inclusive: *const *const bool, } -impl TryFrom<&FFICompactionParams> for CompactionInput { +impl<'a, 'b> TryFrom<&'a FFICompactionParams> for CompactionInput<'a> { type Error = color_eyre::eyre::Report; - fn try_from(params: &FFICompactionParams) -> Result { + fn try_from(params: &'a FFICompactionParams) -> Result, Self::Error> { // We do this separately since we need the values for computing the region let row_key_cols = unpack_string_array(params.row_key_cols, params.row_key_cols_len)? .into_iter() @@ -150,10 +150,10 @@ impl TryFrom<&FFICompactionParams> for CompactionInput { } } -fn compute_region>( - params: &FFICompactionParams, +fn compute_region<'a, T: Borrow>( + params: &'a FFICompactionParams, row_key_cols: &[T], -) -> color_eyre::Result> { +) -> color_eyre::Result>> { let region_mins_inclusive = unpack_primitive_array( params.region_mins_inclusive, params.region_mins_inclusive_len, @@ -360,11 +360,11 @@ fn unpack_primitive_array(array_base: *const *const T, len: usize) -> V /// If the length of the `schema_types` array doesn't match the length specified. /// /// Also panics if a negative array length is found in decoding byte arrays or strings. -fn unpack_variant_array( +fn unpack_variant_array<'a>( array_base: *const *const c_void, len: usize, schema_types: &[i32], -) -> Result, Utf8Error> { +) -> Result>, Utf8Error> { assert_eq!(len, schema_types.len()); unsafe { slice::from_raw_parts(array_base, len) } .iter() diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index b51ee267e4..d6270e4b6d 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -26,9 +26,8 @@ use url::Url; /// Implements a Sleeper compaction algorithm in Rust. /// /// A sequence of Parquet files is read and compacted into a single output Parquet file. The input -/// files must be individually sorted according to the column numbers specified by `sort_columns`. At least -/// one sort column must be specified and one row key column number. A sketches file containing -/// serialised Apache Data Sketches quantiles sketches are written for reach row key column. +/// files must be individually sorted according to the row key columns and then the sort columns`. A sketches file containing +/// serialised Apache Data Sketches quantiles sketches is written for reach row key column. /// #[derive(Parser, Debug)] #[command(author, version)] @@ -50,6 +49,12 @@ struct CmdLineArgs { /// Column names for sort columns #[arg(short = 's', long)] sort_column: Vec, + /// Partition region minimum keys (inclusive). Must be one per row key specified. + #[arg(short='m',long,required=true,num_args=1..)] + region_mins: Vec, + /// Partition region maximum keys (exclusive). Must be one per row key specified. + #[arg(short='n',long,required=true,num_args=1..)] + region_maxs: Vec, } #[tokio::main(flavor = "multi_thread")] @@ -91,16 +96,25 @@ async fn main() -> color_eyre::Result<()> { let output_url = Url::parse(&args.output) .or_else(|_e| Url::parse(&("file://".to_owned() + &args.output)))?; + assert_eq!(args.row_keys.len(), args.region_maxs.len()); + assert_eq!(args.row_keys.len(), args.region_mins.len()); + let mut map = HashMap::new(); - map.insert( - "key".into(), - ColRange { - lower: compaction::PartitionBound::String("h"), - lower_inclusive: true, - upper: compaction::PartitionBound::String("m"), - upper_inclusive: false, - }, - ); + for (key, bounds) in args + .row_keys + .iter() + .zip(args.region_mins.iter().zip(args.region_maxs.iter())) + { + map.insert( + key.into(), + ColRange { + lower: compaction::PartitionBound::String(bounds.0), + lower_inclusive: true, + upper: compaction::PartitionBound::String(bounds.1), + upper_inclusive: false, + }, + ); + } let details = CompactionInput { input_files: input_urls, output_file: output_url, From e06240d4c5df5b8c0eea4354e752e1eede7915e7 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 26 Apr 2024 14:24:35 +0000 Subject: [PATCH 054/129] compiles with sketches --- rust/compaction/src/datafusion.rs | 114 +++++----- .../compaction/src/{ => datafusion}/sketch.rs | 93 ++------ rust/compaction/src/datafusion/udf.rs | 211 ++++++++++++++++++ rust/compaction/src/lib.rs | 1 - 4 files changed, 288 insertions(+), 131 deletions(-) rename rust/compaction/src/{ => datafusion}/sketch.rs (74%) create mode 100644 rust/compaction/src/datafusion/udf.rs diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 14b984ec84..1ec1fe43bc 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -19,24 +19,27 @@ */ use crate::{ aws_s3::{CountingObjectStore, ObjectStoreFactory}, + datafusion::{sketch::serialise_sketches, udf::SketchUDF}, + details::create_sketch_path, ColRange, CompactionInput, CompactionResult, PartitionBound, }; -use arrow::{error::ArrowError, util::pretty::pretty_format_batches}; +use arrow::util::pretty::pretty_format_batches; use datafusion::{ dataframe::DataFrameWriteOptions, error::DataFusionError, execution::{config::SessionConfig, context::SessionContext, options::ParquetReadOptions}, - parquet::{ - arrow::{async_reader::ParquetObjectReader, ParquetRecordBatchStreamBuilder}, - basic::{BrotliLevel, GzipLevel, ZstdLevel}, - }, + logical_expr::ScalarUDF, + parquet::basic::{BrotliLevel, GzipLevel, ZstdLevel}, prelude::*, }; use log::{error, info}; use num_format::{Locale, ToFormattedString}; -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, iter::once, sync::Arc}; use url::Url; +mod sketch; +mod udf; + /// Starts a Sleeper compaction. /// /// The object store factory must be able to produce an [`ObjectStore`] capable of reading @@ -64,24 +67,31 @@ pub async fn compact( let po = ParquetReadOptions::default().file_sort_order(vec![sort_order.clone()]); let mut frame = ctx.read_parquet(input_paths.to_owned(), po).await?; - // Extract all column names - let col_names = frame.schema().clone().strip_qualifiers().field_names(); - info!("All columns in schema {col_names:?}"); - // If we have a partition region, apply it first if let Some(expr) = region_filter(&input_data.region) { frame = frame.filter(expr)?; } + // Create the sketch function + ctx.register_udf(ScalarUDF::from(udf::SketchUDF::new( + frame.schema(), + &input_data.row_key_cols, + ))); + + // Extract all column names + let col_names = frame.schema().clone().strip_qualifiers().field_names(); + info!("All columns in schema {col_names:?}"); + + let row_key_exprs = input_data.row_key_cols.iter().map(col).collect::>(); + let sketch_func = frame.registry().udf("sketch")?; + info!("Using sketch function {sketch_func:?}"); + + let sketch_expr = once(sketch_func.call(row_key_exprs)); // Perform sort of row key and sort columns and projection of all columns - let col_names_expr = frame - .schema() - .clone() - .strip_qualifiers() - .field_names() - .iter() - .map(col) + let col_names_expr = sketch_expr + .chain(col_names.iter().skip(1).map(col)) // 1st column is the sketch function call .collect::>(); + frame = frame.sort(sort_order)?.select(col_names_expr)?; // Show explanation of plan @@ -110,6 +120,41 @@ pub async fn compact( ) .await?; + show_store_stats(store); + + let mut rows_written = 0; + + // Write sketches out to file in Sleeper compatible way + let binding = sketch_func.inner(); + let inner_function: Option<&SketchUDF> = binding.as_any().downcast_ref(); + if let Some(func) = inner_function { + println!( + "Made {} calls to sketch UDF and processed {} total rows.", + func.get_invoke_count(), + func.get_row_count(), + ); + + rows_written = func.get_row_count(); + + // Serialise the sketch + serialise_sketches( + store_factory, + &create_sketch_path(output_path), + &*func.get_sketch(), + ) + .map_err(|e| DataFusionError::External(e.into()))?; + } + + // The rows read will be same as rows_written. + Ok(CompactionResult { + rows_read: rows_written, + rows_written, + }) +} + +/// Show some basic statistics from the [`ObjectStore`]. +/// +fn show_store_stats(store: Arc) { info!( "Object store read {} bytes from {} GETs", store @@ -121,41 +166,6 @@ pub async fn compact( .unwrap_or(0) .to_formatted_string(&Locale::en) ); - - // Find rows just written to output Parquet file - let pq_reader = get_parquet_builder(store_factory, output_path) - .await - .map_err(|e| DataFusionError::External(e.into()))?; - let num_rows = pq_reader.metadata().file_metadata().num_rows(); - - // The rows read will be same as rows_written. - Ok(CompactionResult { - rows_read: num_rows as usize, - rows_written: num_rows as usize, - }) -} - -/// Create an asynchronous builder for reading Parquet files from an object store. -/// -/// The URL must start with a scheme that the object store recognises, e.g. "file" or "s3". -/// -/// # Errors -/// This function will return an error if it couldn't connect to S3 or open a valid -/// Parquet file. -pub async fn get_parquet_builder( - store_factory: &ObjectStoreFactory, - src: &Url, -) -> color_eyre::Result> { - let store = store_factory.get_object_store(src)?; - // HEAD the file to get metadata - let path = object_store::path::Path::from(src.path()); - let object_meta = store - .head(&path) - .await - .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; - // Create a reader for the target file, use it to construct a Stream - let reader = ParquetObjectReader::new(store.as_object_store(), object_meta); - Ok(ParquetRecordBatchStreamBuilder::new(reader).await?) } /// Create the `DataFusion` filtering expression from a Sleeper region. diff --git a/rust/compaction/src/sketch.rs b/rust/compaction/src/datafusion/sketch.rs similarity index 74% rename from rust/compaction/src/sketch.rs rename to rust/compaction/src/datafusion/sketch.rs index cb3c3fd6fb..785af20c23 100644 --- a/rust/compaction/src/sketch.rs +++ b/rust/compaction/src/datafusion/sketch.rs @@ -16,11 +16,8 @@ * limitations under the License. */ use crate::aws_s3::ObjectStoreFactory; -use arrow::array::{ArrayAccessor, AsArray}; -use arrow::datatypes::{ - BinaryType, DataType, Int32Type, Int64Type, LargeBinaryType, LargeUtf8Type, Schema, Utf8Type, -}; -use arrow::record_batch::RecordBatch; +use arrow::array::ArrayAccessor; +use arrow::datatypes::DataType; use bytes::{BufMut, Bytes}; use cxx::{Exception, UniquePtr}; use log::info; @@ -29,14 +26,13 @@ use rust_sketch::quantiles::byte::{byte_sketch_t, new_byte_sketch}; use rust_sketch::quantiles::i32::{i32_sketch_t, new_i32_sketch}; use rust_sketch::quantiles::i64::{i64_sketch_t, new_i64_sketch}; use rust_sketch::quantiles::str::{new_str_sketch, string_sketch_t}; +use std::fmt::Debug; use std::io::Write; -use std::iter::zip; use std::mem::size_of; -use std::sync::Arc; use url::Url; /// Constant size for quantiles data sketches -const K: u16 = 1024; +pub const K: u16 = 1024; pub enum DataSketchVariant { I32(UniquePtr), @@ -45,6 +41,13 @@ pub enum DataSketchVariant { Bytes(DataType, UniquePtr), } +impl Debug for DataSketchVariant { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple(&format!("{}", std::any::type_name::())) + .finish() + } +} + pub trait Item { fn to_i32(&self) -> Option; fn to_i64(&self) -> Option; @@ -160,7 +163,6 @@ impl Item for &[u8] { } } -#[allow(dead_code)] impl DataSketchVariant { /// Updates this sketch variant with the given value. /// @@ -169,7 +171,6 @@ impl DataSketchVariant { /// # Panics /// If the value provided cannot be converted to the type of this [`DataSketch`], i.e. /// if you try to update an i32 sketch with a string. - #[allow(clippy::needless_pass_by_value)] pub fn update(&mut self, value: T) where T: Item, @@ -186,6 +187,7 @@ impl DataSketchVariant { /// /// # Errors /// If the sketch is empty an error is thrown. + #[allow(dead_code)] pub fn get_min_item(&self) -> Result, Exception> { match self { DataSketchVariant::I32(s) => Ok(s.get_min_item().map(Box::new)?), @@ -199,6 +201,7 @@ impl DataSketchVariant { /// /// # Errors /// If the sketch is empty an error is thrown. + #[allow(dead_code)] pub fn get_max_item(&self) -> Result, Exception> { match self { DataSketchVariant::I32(s) => Ok(s.get_max_item().map(Box::new)?), @@ -226,6 +229,7 @@ impl DataSketchVariant { /// Return the underlying Arrow [`DataType`] this sketch is intended for. #[must_use] + #[allow(dead_code)] pub fn data_type(&self) -> DataType { match self { DataSketchVariant::I32(_) => DataType::Int32, @@ -302,61 +306,13 @@ pub fn serialise_sketches( Ok(()) } -/// Update the data sketches for the column numbers in row key fields list. -/// -/// Works through each value in each column indexed by `row_key_fields` and -/// updates the corresponding data sketch i the array of sketches. -/// -/// # Panics -/// Panic if length of sketch array is different to number of row key fields. -/// Panic if `row_key_fields` contains column number that isn't in schema. -/// Panic if there is a mismatch between the sketch data type and the column type. -pub fn update_sketches( - b: &RecordBatch, - sketches: &mut [DataSketchVariant], - row_key_fields: impl AsRef<[usize]>, -) { - assert!( - sketches.len() == row_key_fields.as_ref().len(), - "Differing number of sketches and row key fields!" - ); - // Take each sketch and row key column in turn, then update the sketch - for (sketch, column) in zip( - sketches, - row_key_fields.as_ref().iter().map(|&index| b.column(index)), - ) { - // dynamic dispatch. Match the datatype to the type of sketch to update. - match column.data_type() { - DataType::Int32 => update_sketch(sketch, &column.as_primitive::()), - DataType::Int64 => update_sketch(sketch, &column.as_primitive::()), - DataType::Utf8 => update_sketch( - sketch, - &column.as_string::<::Offset>(), - ), - DataType::LargeUtf8 => update_sketch( - sketch, - &column.as_string::<::Offset>(), - ), - DataType::Binary => update_sketch( - sketch, - &column.as_binary::<::Offset>(), - ), - DataType::LargeBinary => update_sketch( - sketch, - &column.as_binary::<::Offset>(), - ), - _ => panic!("Row type {} not supported", column.data_type()), - } - } -} - /// Update the given sketch from an array. /// /// The list of values in the given array are updated into the data sketch. /// /// # Panics /// Panic if sketch type is not compatible with the item type of the array. -fn update_sketch>( +pub fn update_sketch>( sketch: &mut DataSketchVariant, array: &A, ) { @@ -366,22 +322,3 @@ fn update_sketch>( } } } - -/// Create a vector of Data Sketches. -/// -/// This creates the appropriate data sketch implementations based on on the row key fields -/// and the data types in the schema. Each type is wrapped in a [`SketchEnum`] variant type. -/// -/// # Panics -/// If a row key field can't be found in the schema. -/// -pub fn make_sketches_for_schema( - schema: &Arc, - row_key_fields: &impl AsRef<[usize]>, -) -> Vec { - row_key_fields - .as_ref() - .iter() - .map(|&index| DataSketchVariant::new(schema.field(index).data_type(), K)) - .collect() -} diff --git a/rust/compaction/src/datafusion/udf.rs b/rust/compaction/src/datafusion/udf.rs new file mode 100644 index 0000000000..28ea76bad3 --- /dev/null +++ b/rust/compaction/src/datafusion/udf.rs @@ -0,0 +1,211 @@ +/// A user defined function for `DataFusion` to create quantile sketches +/// of data as it is being compacted. +/// +/// This function is designed to be used with a array of columns, one per Sleeper +/// row key field. The return value is just the first column, untransformed. The sketches +/// are produced as a side effect. +/* +* 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. +*/ +use std::{ + any::Any, + fmt::Debug, + iter::zip, + ops::Deref, + sync::{Mutex, MutexGuard}, +}; + +use arrow::{ + array::AsArray, + datatypes::{ + BinaryType, DataType, Int32Type, Int64Type, LargeBinaryType, LargeUtf8Type, Utf8Type, + }, +}; +use datafusion::{ + common::{internal_err, DFSchema, Result}, + logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility}, + scalar::ScalarValue, +}; + +use super::sketch::{update_sketch, DataSketchVariant, K}; + +pub(crate) struct SketchUDF { + signature: Signature, + invoke_count: Mutex, + row_count: Mutex, + sketch: Mutex>, +} + +impl Debug for SketchUDF { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Foo") + .field("signature", &self.signature) + .field("invoke_count", &self.invoke_count) + .field("row_count", &self.row_count) + .field("sketch", &self.sketch) + .finish() + } +} + +impl SketchUDF { + /// Create a new sketch function based on the schema of the row key fields. + /// + pub fn new(schema: &DFSchema, row_keys: &[String]) -> Self { + Self { + signature: Signature::exact(get_row_key_types(schema, row_keys), Volatility::Immutable), + invoke_count: Mutex::default(), + row_count: Mutex::default(), + sketch: Mutex::new(make_sketches_for_schema(schema, row_keys)), + } + } + + pub fn get_sketch(&self) -> MutexGuard<'_, Vec> { + self.sketch.lock().unwrap() + } + + pub fn get_invoke_count(&self) -> usize { + *self.invoke_count.lock().unwrap() + } + + pub fn get_row_count(&self) -> usize { + *self.row_count.lock().unwrap() + } +} + +/// Create a [`Vec`] of data types for this schema from the row keys. +/// +/// # Panics +/// If a row key field can't be found in the schema. +fn get_row_key_types(schema: &DFSchema, row_keys: &[String]) -> Vec { + row_keys + .iter() + .map(|name| { + schema + .field_with_unqualified_name(name) + .unwrap() + .data_type() + .to_owned() + }) + .collect() +} + +/// Create a vector of Data Sketches. +/// +/// This creates the appropriate data sketch implementations based on on the row key fields +/// and the data types in the schema. Each type is wrapped in a [`SketchEnum`] variant type. +/// +/// # Panics +/// If a row key field can't be found in the schema. +/// +fn make_sketches_for_schema( + schema: &DFSchema, + row_key_fields: &[String], +) -> Vec { + row_key_fields + .iter() + .map(|name| { + DataSketchVariant::new( + schema + .field_with_unqualified_name(name) + .unwrap() + .data_type(), + K, + ) + }) + .collect() +} + +impl ScalarUDFImpl for SketchUDF { + fn as_any(&self) -> &dyn Any { + self + } + fn name(&self) -> &str { + "sketch" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, args: &[DataType]) -> Result { + // Return type will be type of first row key column + Ok(args[0].to_owned()) + } + + fn invoke(&self, columns: &[ColumnarValue]) -> Result { + *self.invoke_count.lock().unwrap() += 1; + + let mut sk_lock = self.sketch.lock().unwrap(); + + for (sketch, col) in zip(sk_lock.iter_mut(), columns) { + match col { + ColumnarValue::Array(array) => { + // dynamic dispatch. Match the datatype to the type of sketch to update. + match array.data_type() { + DataType::Int32 => update_sketch(sketch, &array.as_primitive::()), + DataType::Int64 => update_sketch(sketch, &array.as_primitive::()), + DataType::Utf8 => update_sketch( + sketch, + &array.as_string::<::Offset>(), + ), + DataType::LargeUtf8 => update_sketch( + sketch, + &array.as_string::<::Offset>(), + ), + DataType::Binary => update_sketch( + sketch, + &array.as_binary::<::Offset>(), + ), + DataType::LargeBinary => update_sketch( + sketch, + &array.as_binary::<::Offset>(), + ), + _ => return internal_err!("Row type {} not supported for Sleeper row key field", array.data_type()), + } + *self.row_count.lock().unwrap() += array.len(); + } + + ColumnarValue::Scalar( + ScalarValue::Utf8(Some(value)) | ScalarValue::LargeUtf8(Some(value)), + ) => { + sketch.update(value.deref()); + } + ColumnarValue::Scalar( + ScalarValue::Binary(Some(value)) | ScalarValue::LargeBinary(Some(value)), + ) => { + sketch.update(value.deref()); + } + ColumnarValue::Scalar(ScalarValue::Int32(Some(value))) => { + sketch.update(*value); + } + ColumnarValue::Scalar(ScalarValue::Int64(Some(value))) => { + sketch.update(*value); + } + x @ _ => { + return internal_err!( + "Row type {} not supported for Sleeper row key field", + x.data_type() + ) + } + } + } + + // All columns are same length so lookup row count from first one + *self.row_count.lock().unwrap() += match &columns[0] { + ColumnarValue::Array(array) => array.len(), + ColumnarValue::Scalar(_) => 1, + }; + + Ok(columns[0].clone()) + } +} diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 4cf1de34bb..6a344a1340 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -24,7 +24,6 @@ mod aws_s3; mod datafusion; mod details; -mod sketch; use chrono::Local; use color_eyre::eyre::eyre; From c36c4d313ab312142d249881bb29a2e642497e46 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 26 Apr 2024 14:32:23 +0000 Subject: [PATCH 055/129] Clippy warnings --- rust/compaction/src/datafusion.rs | 21 +++++++++------------ rust/compaction/src/datafusion/sketch.rs | 8 ++++---- rust/compaction/src/datafusion/udf.rs | 13 ++++++------- rust/compaction/src/details.rs | 8 ++++---- rust/compaction/src/lib.rs | 4 ++-- rust/compactor/src/bin/main.rs | 2 +- 6 files changed, 26 insertions(+), 30 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 1ec1fe43bc..8d245340dd 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -120,7 +120,7 @@ pub async fn compact( ) .await?; - show_store_stats(store); + show_store_stats(&store); let mut rows_written = 0; @@ -140,7 +140,7 @@ pub async fn compact( serialise_sketches( store_factory, &create_sketch_path(output_path), - &*func.get_sketch(), + &func.get_sketch(), ) .map_err(|e| DataFusionError::External(e.into()))?; } @@ -154,7 +154,7 @@ pub async fn compact( /// Show some basic statistics from the [`ObjectStore`]. /// -fn show_store_stats(store: Arc) { +fn show_store_stats(store: &Arc) { info!( "Object store read {} bytes from {} GETs", store @@ -194,12 +194,11 @@ fn region_filter(region: &HashMap) -> Option { /// fn upper_bound_expr(range: &ColRange, name: &String) -> Expr { let max_bound = bound_to_lit_expr(&range.upper); - let upper_expr = if range.upper_inclusive { + if range.upper_inclusive { col(name).lt_eq(max_bound) } else { col(name).lt(max_bound) - }; - upper_expr + } } /// Calculate the lower bound expression on a given [`ColRange`]. @@ -208,25 +207,23 @@ fn upper_bound_expr(range: &ColRange, name: &String) -> Expr { /// fn lower_bound_expr(range: &ColRange, name: &String) -> Expr { let min_bound = bound_to_lit_expr(&range.lower); - let lower_expr = if range.lower_inclusive { + if range.lower_inclusive { col(name).gt_eq(min_bound) } else { col(name).gt(min_bound) - }; - lower_expr + } } /// Convert a [`PartitionBound`] to an [`Expr`] that can be /// used in a bigger expression. /// fn bound_to_lit_expr(bound: &PartitionBound) -> Expr { - let expr = match bound { + match bound { PartitionBound::Int32(val) => lit(*val), PartitionBound::Int64(val) => lit(*val), PartitionBound::String(val) => lit(val.to_owned()), PartitionBound::ByteArray(val) => lit(val.to_owned()), - }; - expr + } } /// Convert a Sleeper compression codec string to one `DataFusion` understands. diff --git a/rust/compaction/src/datafusion/sketch.rs b/rust/compaction/src/datafusion/sketch.rs index 785af20c23..89debce5af 100644 --- a/rust/compaction/src/datafusion/sketch.rs +++ b/rust/compaction/src/datafusion/sketch.rs @@ -43,8 +43,7 @@ pub enum DataSketchVariant { impl Debug for DataSketchVariant { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_tuple(&format!("{}", std::any::type_name::())) - .finish() + f.debug_tuple(std::any::type_name::()).finish() } } @@ -171,7 +170,7 @@ impl DataSketchVariant { /// # Panics /// If the value provided cannot be converted to the type of this [`DataSketch`], i.e. /// if you try to update an i32 sketch with a string. - pub fn update(&mut self, value: T) + pub fn update(&mut self, value: &T) where T: Item, { @@ -312,13 +311,14 @@ pub fn serialise_sketches( /// /// # Panics /// Panic if sketch type is not compatible with the item type of the array. +#[allow(clippy::module_name_repetitions)] pub fn update_sketch>( sketch: &mut DataSketchVariant, array: &A, ) { for i in 0..array.len() { unsafe { - sketch.update(array.value_unchecked(i)); + sketch.update(&array.value_unchecked(i)); } } } diff --git a/rust/compaction/src/datafusion/udf.rs b/rust/compaction/src/datafusion/udf.rs index 28ea76bad3..0ff29751fc 100644 --- a/rust/compaction/src/datafusion/udf.rs +++ b/rust/compaction/src/datafusion/udf.rs @@ -23,7 +23,6 @@ use std::{ any::Any, fmt::Debug, iter::zip, - ops::Deref, sync::{Mutex, MutexGuard}, }; @@ -139,7 +138,7 @@ impl ScalarUDFImpl for SketchUDF { } fn return_type(&self, args: &[DataType]) -> Result { // Return type will be type of first row key column - Ok(args[0].to_owned()) + Ok(args[0].clone()) } fn invoke(&self, columns: &[ColumnarValue]) -> Result { @@ -178,20 +177,20 @@ impl ScalarUDFImpl for SketchUDF { ColumnarValue::Scalar( ScalarValue::Utf8(Some(value)) | ScalarValue::LargeUtf8(Some(value)), ) => { - sketch.update(value.deref()); + sketch.update(value); } ColumnarValue::Scalar( ScalarValue::Binary(Some(value)) | ScalarValue::LargeBinary(Some(value)), ) => { - sketch.update(value.deref()); + sketch.update(value); } ColumnarValue::Scalar(ScalarValue::Int32(Some(value))) => { - sketch.update(*value); + sketch.update(value); } ColumnarValue::Scalar(ScalarValue::Int64(Some(value))) => { - sketch.update(*value); + sketch.update(value); } - x @ _ => { + x @ ColumnarValue::Scalar(_) => { return internal_err!( "Row type {} not supported for Sleeper row key field", x.data_type() diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index c37a218dd5..54779023a3 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -54,10 +54,10 @@ pub struct CompactionInput<'a> { impl Default for CompactionInput<'_> { fn default() -> Self { Self { - input_files: Default::default(), + input_files: Vec::default(), output_file: Url::parse("file:///").unwrap(), - row_key_cols: Default::default(), - sort_key_cols: Default::default(), + row_key_cols: Vec::default(), + sort_key_cols: Vec::default(), max_row_group_size: 1_000_000, max_page_size: 65535, compression: "zstd".into(), @@ -67,7 +67,7 @@ impl Default for CompactionInput<'_> { dict_enc_row_keys: true, dict_enc_sort_keys: true, dict_enc_values: true, - region: Default::default(), + region: HashMap::default(), } } } diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 6a344a1340..5ce942a54b 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -107,7 +107,7 @@ pub struct FFICompactionParams { region_maxs_inclusive: *const *const bool, } -impl<'a, 'b> TryFrom<&'a FFICompactionParams> for CompactionInput<'a> { +impl<'a> TryFrom<&'a FFICompactionParams> for CompactionInput<'a> { type Error = color_eyre::eyre::Report; fn try_from(params: &'a FFICompactionParams) -> Result, Self::Error> { @@ -387,7 +387,7 @@ fn unpack_variant_array<'a>( #[allow(clippy::cast_sign_loss)] slice::from_raw_parts(bptr.byte_add(4).cast::(), str_len as usize) }) - .map(|v| PartitionBound::String(v)) + .map(PartitionBound::String) } 4 => { //unpack length (signed because it's from Java) diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index d6270e4b6d..5ce11695fd 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -26,7 +26,7 @@ use url::Url; /// Implements a Sleeper compaction algorithm in Rust. /// /// A sequence of Parquet files is read and compacted into a single output Parquet file. The input -/// files must be individually sorted according to the row key columns and then the sort columns`. A sketches file containing +/// files must be individually sorted according to the row key columns and then the sort columns. A sketches file containing /// serialised Apache Data Sketches quantiles sketches is written for reach row key column. /// #[derive(Parser, Debug)] From 7d84edcee89adc1fdd1adc9e0b657013bb4d5ea3 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 26 Apr 2024 15:13:31 +0000 Subject: [PATCH 056/129] Sketches work --- rust/compaction/src/datafusion.rs | 24 +++++++++++++++++------- rust/compactor/src/bin/main.rs | 20 +++++++++++++------- 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 8d245340dd..413c031a0b 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -27,7 +27,10 @@ use arrow::util::pretty::pretty_format_batches; use datafusion::{ dataframe::DataFrameWriteOptions, error::DataFusionError, - execution::{config::SessionConfig, context::SessionContext, options::ParquetReadOptions}, + execution::{ + config::SessionConfig, context::SessionContext, options::ParquetReadOptions, + FunctionRegistry, + }, logical_expr::ScalarUDF, parquet::basic::{BrotliLevel, GzipLevel, ZstdLevel}, prelude::*, @@ -73,25 +76,32 @@ pub async fn compact( } // Create the sketch function - ctx.register_udf(ScalarUDF::from(udf::SketchUDF::new( + let sketch_func = Arc::new(ScalarUDF::from(udf::SketchUDF::new( frame.schema(), &input_data.row_key_cols, ))); + frame.task_ctx().register_udf(sketch_func.clone())?; // Extract all column names let col_names = frame.schema().clone().strip_qualifiers().field_names(); info!("All columns in schema {col_names:?}"); let row_key_exprs = input_data.row_key_cols.iter().map(col).collect::>(); - let sketch_func = frame.registry().udf("sketch")?; info!("Using sketch function {sketch_func:?}"); - let sketch_expr = once(sketch_func.call(row_key_exprs)); + let sketch_expr = once( + sketch_func + .call(row_key_exprs) + .alias(&input_data.row_key_cols[0]), + ); // Perform sort of row key and sort columns and projection of all columns let col_names_expr = sketch_expr .chain(col_names.iter().skip(1).map(col)) // 1st column is the sketch function call .collect::>(); + let _no_sketches = col_names.iter().map(col).collect::>(); + + // Build compaction query frame = frame.sort(sort_order)?.select(col_names_expr)?; // Show explanation of plan @@ -128,10 +138,10 @@ pub async fn compact( let binding = sketch_func.inner(); let inner_function: Option<&SketchUDF> = binding.as_any().downcast_ref(); if let Some(func) = inner_function { - println!( + info!( "Made {} calls to sketch UDF and processed {} total rows.", - func.get_invoke_count(), - func.get_row_count(), + func.get_invoke_count().to_formatted_string(&Locale::en), + func.get_row_count().to_formatted_string(&Locale::en), ); rows_written = func.get_row_count(); diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index 5ce11695fd..f1d20c3d1c 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -16,6 +16,7 @@ use chrono::Local; use clap::Parser; +use color_eyre::eyre::bail; use compaction::{merge_sorted_files, ColRange, CompactionInput}; use human_panic::setup_panic; use log::info; @@ -133,12 +134,17 @@ async fn main() -> color_eyre::Result<()> { }; let result = merge_sorted_files(&details).await; - if let Ok(data) = result { - info!( - "Compaction read {} rows and wrote {} rows", - data.rows_read.to_formatted_string(&Locale::en), - data.rows_written.to_formatted_string(&Locale::en) - ); - } + match result { + Ok(r) => { + info!( + "Compaction read {} rows and wrote {} rows", + r.rows_read.to_formatted_string(&Locale::en), + r.rows_written.to_formatted_string(&Locale::en) + ); + } + Err(e) => { + bail!(e); + } + }; Ok(()) } From 6484759a3aa6c3f4311a6bba69ffdd14768d18a5 Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 30 Apr 2024 08:41:40 +0000 Subject: [PATCH 057/129] Fixed log messages --- .../job/execution/DefaultSelector.java | 4 +-- java/compaction/compaction-rust/pom.xml | 13 +++------ rust/Cargo.lock | 28 +++++++++---------- 3 files changed, 20 insertions(+), 25 deletions(-) diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java index bf77b1266c..ba3ebafc9c 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java @@ -75,11 +75,11 @@ public CompactionRunner chooseCompactor(CompactionJob job) { // Is an iterator specifed, if so can we support this? if (job.getIteratorClassName() != null && !runner.supportsIterators()) { - LOGGER.debug("Table has an iterator set, which compactor %s doesn't support, falling back to default", runner.getClass().getSimpleName()); + LOGGER.debug("Table has an iterator set, which compactor {} doesn't support, falling back to default", runner.getClass().getSimpleName()); runner = defaultRunner; } - LOGGER.info("Selecting %s compactor (language %s) for job ID %s table ID %s", runner.getClass().getSimpleName(), runner.implementationLanguage(), job.getId(), job.getTableId()); + LOGGER.info("Selecting {} compactor (language {}) for job ID {} table ID {}", runner.getClass().getSimpleName(), runner.implementationLanguage(), job.getId(), job.getTableId()); return runner; } } diff --git a/java/compaction/compaction-rust/pom.xml b/java/compaction/compaction-rust/pom.xml index 34dbf0f624..69f01d8fb9 100644 --- a/java/compaction/compaction-rust/pom.xml +++ b/java/compaction/compaction-rust/pom.xml @@ -105,11 +105,6 @@ - - - org.apache.maven.plugins - maven-shade-plugin - org.apache.maven.plugins @@ -180,7 +175,7 @@ - + diff --git a/rust/Cargo.lock b/rust/Cargo.lock index e070849a16..6bedbf15b3 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1567,9 +1567,9 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.0.2" +version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "658bd65b1cf4c852a3cc96f18a8ce7b5640f6b703f905c7d74532294c2a63984" +checksum = "9fc0510504f03c51ada170672ac806f1f105a88aa97a5281117e1ddc3368e51a" [[package]] name = "fixedbitset" @@ -1796,9 +1796,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.14.3" +version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "290f1a1d9242c78d09ce40a5e87e7554ee637af1351968159f4952f028f75604" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash", "allocator-api2", @@ -2885,9 +2885,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.21.11" +version = "0.21.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fecbfb7b1444f477b345853b1fce097a2c6fb637b2bfb87e6bc5db0f043fae4" +checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" dependencies = [ "log", "ring", @@ -3015,18 +3015,18 @@ checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" [[package]] name = "serde" -version = "1.0.198" +version = "1.0.199" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9846a40c979031340571da2545a4e5b7c4163bdae79b301d5f86d03979451fcc" +checksum = "0c9f6e76df036c77cd94996771fb40db98187f096dd0b9af39c6c6e452ba966a" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.198" +version = "1.0.199" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e88edab869b01783ba905e7d0153f9fc1a6505a96e4ad3018011eedb838566d9" +checksum = "11bd257a6541e141e42ca6d24ae26f7714887b47e89aa739099104c7e4d3b7fc" dependencies = [ "proc-macro2", "quote", @@ -3145,9 +3145,9 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" [[package]] name = "socket2" -version = "0.5.6" +version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05ffd9c0a93b7543e062e759284fcf5f5e3b098501104bfbdde4d404db792871" +checksum = "ce305eb0b4296696835b71df73eb912e0f1ffd2556a501fcede6e0c50349191c" dependencies = [ "libc", "windows-sys 0.52.0", @@ -3606,9 +3606,9 @@ checksum = "d4c87d22b6e3f4a18d4d40ef354e97c90fcb14dd91d7dc0aa9d8a1172ebf7202" [[package]] name = "unicode-width" -version = "0.1.11" +version = "0.1.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" +checksum = "68f5e5f3158ecfd4b8ff6fe086db7c8467a2dfdac97fe420f2b7c4aa97af66d6" [[package]] name = "untrusted" From 8196b0ff9dbc68cc1c4d85efcd4a07d96452189c Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 30 Apr 2024 11:23:13 +0000 Subject: [PATCH 058/129] Null upper bounds now supported --- .../job/CompactionRunnerDetails.java | 16 ++++ .../job/execution/DefaultSelector.java | 3 +- .../execution}/RustBridge.java | 21 ++--- .../execution}/RustCompaction.java | 39 +++++++--- rust/compaction/src/datafusion.rs | 61 ++++++++++----- rust/compaction/src/details.rs | 2 + rust/compaction/src/lib.rs | 77 ++++++++++++------- rust/compactor/src/bin/main.rs | 8 +- 8 files changed, 157 insertions(+), 70 deletions(-) rename java/compaction/compaction-rust/src/main/java/sleeper/compaction/{jobexecution => job/execution}/RustBridge.java (96%) rename java/compaction/compaction-rust/src/main/java/sleeper/compaction/{jobexecution => job/execution}/RustCompaction.java (92%) diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/CompactionRunnerDetails.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/CompactionRunnerDetails.java index 0d24e8f933..d856933fd4 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/CompactionRunnerDetails.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/CompactionRunnerDetails.java @@ -16,14 +16,30 @@ package sleeper.compaction.job; public interface CompactionRunnerDetails { + /** + * Some compaction implementations may use hardware acceleration such as GPUs. + * + * @return true iff this compaction implementation uses any sort of hardware acceleration + */ default boolean isHardwareAccelerated() { return false; } + /** + * What language is this implemented in? If multiple languages are used, the primary + * one used for performing the compaction computation should be returned. + * + * @return the principal implementation language for this compactor + */ default String implementationLanguage() { return "Java"; } + /** + * States whether this compactor compact Sleeper tables that have iterators attached to them. + * + * @return true if iterators can be processed by this compactor + */ default boolean supportsIterators() { return false; } diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java index ba3ebafc9c..39be52c8f2 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java @@ -20,7 +20,6 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionRunner; -import sleeper.compaction.jobexecution.RustCompaction; import sleeper.configuration.jars.ObjectFactory; import sleeper.configuration.properties.instance.InstanceProperties; import sleeper.configuration.properties.table.TableProperties; @@ -73,7 +72,7 @@ public CompactionRunner chooseCompactor(CompactionJob job) { break; } - // Is an iterator specifed, if so can we support this? + // Is an iterator specifed? If so can we support this? if (job.getIteratorClassName() != null && !runner.supportsIterators()) { LOGGER.debug("Table has an iterator set, which compactor {} doesn't support, falling back to default", runner.getClass().getSimpleName()); runner = defaultRunner; diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/job/execution/RustBridge.java similarity index 96% rename from java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java rename to java/compaction/compaction-rust/src/main/java/sleeper/compaction/job/execution/RustBridge.java index 971e47d33f..080a1731ec 100644 --- a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustBridge.java +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/job/execution/RustBridge.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.jobexecution; +package sleeper.compaction.job.execution; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import jnr.ffi.LibraryLoader; @@ -238,22 +238,26 @@ public Array(Struct enclosing) { * A base pointer is allocated pointers set to other * dynamically allocated memory containing items from array. * - * @param arr array data + * @param arr array data + * @param nullsAllowed if null pointers are allowed in the data array */ - public void populate(final T[] arr) { + public void populate(final T[] arr, boolean nullsAllowed) { final jnr.ffi.Runtime r = len.struct().getRuntime(); // Calculate size needed for array of pointers int ptrSize = r.findType(NativeType.ADDRESS).size(); // Null out zero length arrays if (arr.length > 0) { int size = arr.length * ptrSize; - // Allocate some memory for string pointers + // Allocate some memory for pointers this.basePtr = r.getMemoryManager().allocateDirect(size); this.arrayBase.set(basePtr); this.items = new jnr.ffi.Pointer[arr.length]; for (int i = 0; i < arr.length; i++) { + if (!nullsAllowed && arr[i] == null) { + throw new NullPointerException("Index " + i + " of array is null when nulls aren't allowed here"); + } setValue(arr[i], i, r); } @@ -287,9 +291,6 @@ public void validate() { if (this.arrayBase.get().address() != this.basePtr.address()) { throw new IllegalStateException("array base pointer and stored base pointer differ!"); } - for (jnr.ffi.Pointer p : items) { - Objects.requireNonNull(p, "NULL pointer found in pointer array store"); - } } } @@ -307,7 +308,9 @@ public void validate() { * @throws IndexOutOfBoundsException if idx is invalid */ protected void setValue(E item, int idx, jnr.ffi.Runtime r) { - if (item instanceof Integer) { + if (item == null) { + this.items[idx] = jnr.ffi.Pointer.wrap(r, 0); + } else if (item instanceof Integer) { int e = (int) item; this.items[idx] = r.getMemoryManager().allocateDirect(r.findType(NativeType.SINT).size()); this.items[idx].putInt(0, e); @@ -321,7 +324,7 @@ protected void setValue(E item, int idx, jnr.ffi.Runtime r) { byte[] utf8string = e.getBytes(StandardCharsets.UTF_8); // Add four for length int stringSize = utf8string.length + 4; - // Allocate memory for string and write length then string + // Allocate memory for string and write length then the string this.items[idx] = r.getMemoryManager().allocateDirect(stringSize); this.items[idx].putInt(0, utf8string.length); this.items[idx].put(4, utf8string, 0, utf8string.length); diff --git a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/job/execution/RustCompaction.java similarity index 92% rename from java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java rename to java/compaction/compaction-rust/src/main/java/sleeper/compaction/job/execution/RustCompaction.java index ed337a89af..ac29cc6542 100644 --- a/java/compaction/compaction-rust/src/main/java/sleeper/compaction/jobexecution/RustCompaction.java +++ b/java/compaction/compaction-rust/src/main/java/sleeper/compaction/job/execution/RustCompaction.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.jobexecution; +package sleeper.compaction.job.execution; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -21,7 +21,7 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionRunner; import sleeper.compaction.job.StateStoreUpdate; -import sleeper.compaction.jobexecution.RustBridge.FFICompactionParams; +import sleeper.compaction.job.execution.RustBridge.FFICompactionParams; import sleeper.configuration.properties.table.TableProperties; import sleeper.configuration.properties.table.TablePropertiesProvider; import sleeper.core.range.Range; @@ -108,11 +108,11 @@ public RecordsProcessed compact(CompactionJob job) throws Exception { @SuppressWarnings(value = "checkstyle:avoidNestedBlocks") public static FFICompactionParams createFFIParams(CompactionJob job, TableProperties tableProperties, Schema schema, Region region, jnr.ffi.Runtime runtime) { FFICompactionParams params = new FFICompactionParams(runtime); - params.input_files.populate(job.getInputFiles().toArray(new String[0])); + params.input_files.populate(job.getInputFiles().toArray(new String[0]), false); params.output_file.set(job.getOutputFile()); - params.row_key_cols.populate(schema.getRowKeyFieldNames().toArray(new String[0])); - params.row_key_schema.populate(getKeyTypes(schema.getRowKeyTypes())); - params.sort_key_cols.populate(schema.getSortKeyFieldNames().toArray(new String[0])); + params.row_key_cols.populate(schema.getRowKeyFieldNames().toArray(new String[0]), false); + params.row_key_schema.populate(getKeyTypes(schema.getRowKeyTypes()), false); + params.sort_key_cols.populate(schema.getSortKeyFieldNames().toArray(new String[0]), false); params.max_row_group_size.set(RUST_MAX_ROW_GROUP_ROWS); params.max_page_size.set(tableProperties.getInt(PAGE_SIZE)); params.compression.set(tableProperties.get(COMPRESSION_CODEC)); @@ -122,22 +122,24 @@ public static FFICompactionParams createFFIParams(CompactionJob job, TableProper params.dict_enc_row_keys.set(tableProperties.getBoolean(DICTIONARY_ENCODING_FOR_ROW_KEY_FIELDS)); params.dict_enc_sort_keys.set(tableProperties.getBoolean(DICTIONARY_ENCODING_FOR_SORT_KEY_FIELDS)); params.dict_enc_values.set(tableProperties.getBoolean(DICTIONARY_ENCODING_FOR_VALUE_FIELDS)); - // Sanity check: minimise lifetime + // Extra braces: Make sure wrong array isn't populated to wrong pointers { + // This array can't contain nulls Object[] regionMins = region.getRanges().stream().map(Range::getMin).toArray(); - params.region_mins.populate(regionMins); + params.region_mins.populate(regionMins, false); } { Boolean[] regionMinInclusives = region.getRanges().stream().map(Range::isMinInclusive).toArray(Boolean[]::new); - params.region_mins_inclusive.populate(regionMinInclusives); + params.region_mins_inclusive.populate(regionMinInclusives, false); } { + // This array can contain nulls Object[] regionMaxs = region.getRanges().stream().map(Range::getMax).toArray(); - params.region_maxs.populate(regionMaxs); + params.region_maxs.populate(regionMaxs, true); } { Boolean[] regionMaxInclusives = region.getRanges().stream().map(Range::isMaxInclusive).toArray(Boolean[]::new); - params.region_maxs_inclusive.populate(regionMaxInclusives); + params.region_maxs_inclusive.populate(regionMaxInclusives, false); } params.validate(); return params; @@ -204,4 +206,19 @@ public static RecordsProcessed invokeRustFFI(CompactionJob job, RustBridge.Compa nativeLib.free_result(compactionData); } } + + @Override + public String implementationLanguage() { + return "Rust"; + } + + @Override + public boolean isHardwareAccelerated() { + return false; + } + + @Override + public boolean supportsIterators() { + return false; + } } diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 413c031a0b..8fec27031e 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -62,9 +62,9 @@ pub async fn compact( // Register some object store from first input file and output file let store = register_store(store_factory, input_paths, output_path, &ctx)?; - // Sort on row key columns then sort columns (nulls last) + // Sort on row key columns then sort key columns (nulls last) let sort_order = sort_order(input_data); - info!("Row key and sort column order {sort_order:?}"); + info!("Row key and sort key column order {sort_order:?}"); // Tell DataFusion that the row key columns and sort columns are already sorted let po = ParquetReadOptions::default().file_sort_order(vec![sort_order.clone()]); @@ -94,7 +94,7 @@ pub async fn compact( .call(row_key_exprs) .alias(&input_data.row_key_cols[0]), ); - // Perform sort of row key and sort columns and projection of all columns + // Perform sort of row key and sort key columns and projection of all columns let col_names_expr = sketch_expr .chain(col_names.iter().skip(1).map(col)) // 1st column is the sketch function call .collect::>(); @@ -188,11 +188,18 @@ fn region_filter(region: &HashMap) -> Option { for (name, range) in region { let lower_expr = lower_bound_expr(range, name); let upper_expr = upper_bound_expr(range, name); - let expr = lower_expr.and(upper_expr); + let expr = match (lower_expr, upper_expr) { + (Some(l), Some(u)) => Some(l.and(u)), + (Some(l), None) => Some(l), + (None, Some(u)) => Some(u), + (None, None) => None, + }; // Combine this column filter with any previous column filter - col_expr = match col_expr { - Some(original) => Some(original.and(expr)), - None => Some(expr), + if let Some(e) = expr { + col_expr = match col_expr { + Some(original) => Some(original.and(e)), + None => Some(e), + } } } col_expr @@ -202,37 +209,55 @@ fn region_filter(region: &HashMap) -> Option { /// /// This takes into account the inclusive/exclusive nature of the bound. /// -fn upper_bound_expr(range: &ColRange, name: &String) -> Expr { - let max_bound = bound_to_lit_expr(&range.upper); - if range.upper_inclusive { - col(name).lt_eq(max_bound) +fn upper_bound_expr(range: &ColRange, name: &String) -> Option { + if let PartitionBound::Unbounded = range.upper { + None } else { - col(name).lt(max_bound) + let max_bound = bound_to_lit_expr(&range.upper); + if range.upper_inclusive { + Some(col(name).lt_eq(max_bound)) + } else { + Some(col(name).lt(max_bound)) + } } } /// Calculate the lower bound expression on a given [`ColRange`]. /// +/// Not all bounds are present, so `None` is returned for the unbounded case. +/// /// This takes into account the inclusive/exclusive nature of the bound. /// -fn lower_bound_expr(range: &ColRange, name: &String) -> Expr { - let min_bound = bound_to_lit_expr(&range.lower); - if range.lower_inclusive { - col(name).gt_eq(min_bound) +fn lower_bound_expr(range: &ColRange, name: &String) -> Option { + if let PartitionBound::Unbounded = range.lower { + None } else { - col(name).gt(min_bound) + let min_bound = bound_to_lit_expr(&range.lower); + if range.lower_inclusive { + Some(col(name).gt_eq(min_bound)) + } else { + Some(col(name).gt(min_bound)) + } } } /// Convert a [`PartitionBound`] to an [`Expr`] that can be /// used in a bigger expression. /// +/// # Panics +/// If bound is [`PartitionBound::Unbounded`] as we can't construct +/// an expression for that. +/// fn bound_to_lit_expr(bound: &PartitionBound) -> Expr { match bound { PartitionBound::Int32(val) => lit(*val), PartitionBound::Int64(val) => lit(*val), PartitionBound::String(val) => lit(val.to_owned()), PartitionBound::ByteArray(val) => lit(val.to_owned()), + PartitionBound::Unbounded => { + error!("Can't create filter expression for unbounded partition range!"); + panic!("Can't create filter expression for unbounded partition range!"); + } } } @@ -270,7 +295,7 @@ fn create_session_cfg(input_data: &CompactionInput) -> SessionConfig { /// Creates the sort order for a given schema. /// -/// This is a list of the row key columns followed by the sort columns. +/// This is a list of the row key columns followed by the sort key columns. /// fn sort_order(input_data: &CompactionInput) -> Vec { let sort_order = input_data diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 54779023a3..7216648802 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -30,6 +30,8 @@ pub enum PartitionBound<'a> { Int64(i64), String(&'a str), ByteArray(&'a [u8]), + /// Represented by a NULL in Java + Unbounded, } /// All the information for a a Sleeper compaction. diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 5ce942a54b..279b33432d 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -100,6 +100,7 @@ pub struct FFICompactionParams { region_mins_len: usize, region_mins: *const *const c_void, region_maxs_len: usize, + // The region_maxs array may contain null pointers!! region_maxs: *const *const c_void, region_mins_inclusive_len: usize, region_mins_inclusive: *const *const bool, @@ -162,10 +163,18 @@ fn compute_region<'a, T: Borrow>( params.region_maxs_inclusive_len, ); let schema_types = unpack_primitive_array(params.row_key_schema, params.row_key_schema_len); - let region_mins = - unpack_variant_array(params.region_mins, params.region_mins_len, &schema_types)?; - let region_maxs = - unpack_variant_array(params.region_maxs, params.region_maxs_len, &schema_types)?; + let region_mins = unpack_variant_array( + params.region_mins, + params.region_mins_len, + &schema_types, + false, + )?; + let region_maxs = unpack_variant_array( + params.region_maxs, + params.region_maxs_len, + &schema_types, + true, + )?; let mut map = HashMap::with_capacity(row_key_cols.len()); for (idx, row_key) in row_key_cols.iter().enumerate() { @@ -357,49 +366,65 @@ fn unpack_primitive_array(array_base: *const *const T, len: usize) -> V /// /// # Panics /// If the length of the `schema_types` array doesn't match the length specified. +/// If `nulls_present` is false and a null pointer is found. /// /// Also panics if a negative array length is found in decoding byte arrays or strings. fn unpack_variant_array<'a>( array_base: *const *const c_void, len: usize, schema_types: &[i32], + nulls_present: bool, ) -> Result>, Utf8Error> { assert_eq!(len, schema_types.len()); unsafe { slice::from_raw_parts(array_base, len) } .iter() .inspect(|p| { - if p.is_null() { + if !nulls_present && p.is_null() { error!("Found NULL pointer in string array"); } }) .zip(schema_types.iter()) .map(|(&bptr, type_id)| match type_id { - 1 => Ok(PartitionBound::Int32(unsafe { *bptr.cast::() })), - 2 => Ok(PartitionBound::Int64(unsafe { *bptr.cast::() })), + 1 => Ok(match unsafe { bptr.cast::().as_ref() } { + Some(v) => PartitionBound::Int32(*v), + None => PartitionBound::Unbounded, + }), + 2 => Ok(match unsafe { bptr.cast::().as_ref() } { + Some(v) => PartitionBound::Int64(*v), + None => PartitionBound::Unbounded, + }), 3 => { - //unpack length (signed because it's from Java) - let str_len = unsafe { *bptr.cast::() }; - if str_len < 0 { - error!("Illegal string length in FFI array: {str_len}"); - panic!("Illegal string length in FFI array: {str_len}"); + match unsafe { bptr.cast::().as_ref() } { + //unpack length (signed because it's from Java) + Some(str_len) => { + if *str_len < 0 { + error!("Illegal string length in FFI array: {str_len}"); + panic!("Illegal string length in FFI array: {str_len}"); + } + std::str::from_utf8(unsafe { + #[allow(clippy::cast_sign_loss)] + slice::from_raw_parts(bptr.byte_add(4).cast::(), *str_len as usize) + }) + .map(PartitionBound::String) + } + None => Ok(PartitionBound::Unbounded), } - std::str::from_utf8(unsafe { - #[allow(clippy::cast_sign_loss)] - slice::from_raw_parts(bptr.byte_add(4).cast::(), str_len as usize) - }) - .map(PartitionBound::String) } 4 => { - //unpack length (signed because it's from Java) - let byte_len = unsafe { *bptr.cast::() }; - if byte_len < 0 { - error!("Illegal byte array length in FFI array: {byte_len}"); - panic!("Illegal byte array length in FFI array: {byte_len}"); + match unsafe { bptr.cast::().as_ref() } { + //unpack length (signed because it's from Java) + Some(byte_len) => { + if *byte_len < 0 { + error!("Illegal byte array length in FFI array: {byte_len}"); + panic!("Illegal byte array length in FFI array: {byte_len}"); + } + Ok(PartitionBound::ByteArray(unsafe { + #[allow(clippy::cast_sign_loss)] + slice::from_raw_parts(bptr.byte_add(4).cast::(), *byte_len as usize) + })) + } + None => Ok(PartitionBound::Unbounded), } - Ok(PartitionBound::ByteArray(unsafe { - #[allow(clippy::cast_sign_loss)] - slice::from_raw_parts(bptr.byte_add(4).cast::(), byte_len as usize) - })) } x => { error!("Unexpected type id {x}"); diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index f1d20c3d1c..2e81239912 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -44,12 +44,12 @@ struct CmdLineArgs { /// List of input Parquet files (must be sorted) as URLs #[arg(num_args=1.., required=true)] input: Vec, - /// Column names for a row key fields + /// Column names for a row key columns #[arg(short = 'k', long, num_args=1.., required=true)] row_keys: Vec, - /// Column names for sort columns + /// Column names for sort key columns #[arg(short = 's', long)] - sort_column: Vec, + sort_keys: Vec, /// Partition region minimum keys (inclusive). Must be one per row key specified. #[arg(short='m',long,required=true,num_args=1..)] region_mins: Vec, @@ -130,7 +130,7 @@ async fn main() -> color_eyre::Result<()> { dict_enc_values: true, region: map, row_key_cols: args.row_keys, - sort_key_cols: args.sort_column, + sort_key_cols: args.sort_keys, }; let result = merge_sorted_files(&details).await; From c022773e8a0cc78bbe663479e747ce459785da80 Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 30 Apr 2024 11:47:22 +0000 Subject: [PATCH 059/129] Update name of UDF --- rust/compaction/src/datafusion/udf.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/compaction/src/datafusion/udf.rs b/rust/compaction/src/datafusion/udf.rs index 0ff29751fc..fb275e80fb 100644 --- a/rust/compaction/src/datafusion/udf.rs +++ b/rust/compaction/src/datafusion/udf.rs @@ -49,7 +49,7 @@ pub(crate) struct SketchUDF { impl Debug for SketchUDF { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Foo") + f.debug_struct("SketchUDF") .field("signature", &self.signature) .field("invoke_count", &self.invoke_count) .field("row_count", &self.row_count) From c8f75826e511513f29fb337d8f990f0c119d67a2 Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 30 Apr 2024 15:14:04 +0000 Subject: [PATCH 060/129] Some changes in preparation for memory limiting --- rust/compaction/src/datafusion.rs | 29 ++++++++++++++++++++++++----- rust/compactor/src/bin/main.rs | 2 +- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 8fec27031e..9569417954 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -28,7 +28,10 @@ use datafusion::{ dataframe::DataFrameWriteOptions, error::DataFusionError, execution::{ - config::SessionConfig, context::SessionContext, options::ParquetReadOptions, + config::SessionConfig, + context::SessionContext, + options::ParquetReadOptions, + runtime_env::{RuntimeConfig, RuntimeEnv}, FunctionRegistry, }, logical_expr::ScalarUDF, @@ -57,6 +60,9 @@ pub async fn compact( info!("DataFusion compaction of {input_paths:?}"); info!("Compaction region {:?}", input_data.region); let sf = create_session_cfg(input_data); + // let lim = std::env::var("LIM").unwrap().parse::().unwrap(); + // let rt = RuntimeConfig::new().with_memory_limit(lim, 0.8f64); + // let ctx = SessionContext::new_with_config_rt(sf, Arc::new(RuntimeEnv::new(rt)?)); let ctx = SessionContext::new_with_config(sf); // Register some object store from first input file and output file @@ -268,9 +274,21 @@ fn get_compression(compression: &str) -> String { "gzip" => format!("gzip({})", GzipLevel::default().compression_level()), "brotli" => format!("brotli({})", BrotliLevel::default().compression_level()), "zstd" => format!("zstd({})", ZstdLevel::default().compression_level()), - _ => { - error!("Unknown compression"); - unimplemented!() + x => { + error!("Unknown compression {x}, valid values: uncompressed, snappy, lzo, lz4, gzip, brotli, zstd"); + unimplemented!("Unknown compression {x}, valid values: uncompressed, snappy, lzo, lz4, gzip, brotli, zstd"); + } + } +} + +/// Convert a Sleeper Parquet version to one `DataFusion` understands. +fn get_parquet_writer_version(version: &str) -> String { + match version { + "v1" => "1.0".into(), + "v2" => "2.0".into(), + x => { + error!("Parquet writer version invalid {x}, valid values: v1, v2"); + unimplemented!("Parquet writer version invalid {x}, valid values: v1, v2"); } } } @@ -284,7 +302,8 @@ fn create_session_cfg(input_data: &CompactionInput) -> SessionConfig { sf.options_mut().execution.parquet.max_row_group_size = input_data.max_row_group_size; sf.options_mut().execution.parquet.data_pagesize_limit = input_data.max_page_size; sf.options_mut().execution.parquet.compression = Some(get_compression(&input_data.compression)); - sf.options_mut().execution.parquet.writer_version = input_data.writer_version.clone(); + sf.options_mut().execution.parquet.writer_version = + get_parquet_writer_version(&input_data.writer_version); sf.options_mut() .execution .parquet diff --git a/rust/compactor/src/bin/main.rs b/rust/compactor/src/bin/main.rs index 2e81239912..7b8423c73d 100644 --- a/rust/compactor/src/bin/main.rs +++ b/rust/compactor/src/bin/main.rs @@ -124,7 +124,7 @@ async fn main() -> color_eyre::Result<()> { column_truncate_length: 1_048_576, stats_truncate_length: 1_048_576, compression: "zstd".into(), - writer_version: "2.0".into(), + writer_version: "v2".into(), dict_enc_row_keys: true, dict_enc_sort_keys: true, dict_enc_values: true, From 82118d76f7bf94fc2132e8475ee1e56357236aeb Mon Sep 17 00:00:00 2001 From: m09526 Date: Tue, 30 Apr 2024 15:47:20 +0000 Subject: [PATCH 061/129] Write CPU numbers out --- rust/Cargo.lock | 1 + rust/compaction/Cargo.toml | 1 + rust/compaction/src/details.rs | 6 ++++++ 3 files changed, 8 insertions(+) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 6bedbf15b3..22b2cd9446 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1032,6 +1032,7 @@ dependencies = [ "libc", "log", "num-format", + "num_cpus", "object_store", "rust_sketch", "tokio", diff --git a/rust/compaction/Cargo.toml b/rust/compaction/Cargo.toml index f0a7392895..2371818ad9 100644 --- a/rust/compaction/Cargo.toml +++ b/rust/compaction/Cargo.toml @@ -53,3 +53,4 @@ tokio-test = { version = "0.4.2" } # Doc tests env_logger = { version = "0.11.3" } chrono = { version = "0.4.26" } # Log helper color-eyre = { version = "0.6.3" } # Error handling +num_cpus = { version = "1.16.0" } # Read CPU count diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 7216648802..40b93c8f58 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -20,6 +20,7 @@ use aws_config::BehaviorVersion; use aws_credential_types::provider::ProvideCredentials; use color_eyre::eyre::{eyre, Result}; +use log::info; use std::{collections::HashMap, path::PathBuf}; use url::Url; @@ -129,6 +130,11 @@ pub struct CompactionResult { /// There must be at least one input file. /// pub async fn merge_sorted_files(input_data: &CompactionInput<'_>) -> Result { + info!( + "Environment has {} CPUs available {} cores", + num_cpus::get_physical(), + num_cpus::get() + ); // Read the schema from the first file if input_data.input_files.is_empty() { Err(eyre!("No input paths supplied")) From 13fc8218647579cdf9ef18942830119243039eb6 Mon Sep 17 00:00:00 2001 From: m09526 Date: Thu, 2 May 2024 08:51:59 +0000 Subject: [PATCH 062/129] Memory limit expr. --- rust/compaction/src/datafusion.rs | 21 +++++++++++++++++---- rust/compaction/src/details.rs | 22 ++++++++++++---------- 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 9569417954..ad7faf4597 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -30,6 +30,7 @@ use datafusion::{ execution::{ config::SessionConfig, context::SessionContext, + memory_pool::FairSpillPool, options::ParquetReadOptions, runtime_env::{RuntimeConfig, RuntimeEnv}, FunctionRegistry, @@ -59,11 +60,23 @@ pub async fn compact( ) -> Result { info!("DataFusion compaction of {input_paths:?}"); info!("Compaction region {:?}", input_data.region); - let sf = create_session_cfg(input_data); - // let lim = std::env::var("LIM").unwrap().parse::().unwrap(); - // let rt = RuntimeConfig::new().with_memory_limit(lim, 0.8f64); - // let ctx = SessionContext::new_with_config_rt(sf, Arc::new(RuntimeEnv::new(rt)?)); + + let lim = std::env::var("LIM") + .unwrap_or("104857600".into()) + .parse::() + .unwrap(); + + let mut sf = create_session_cfg(input_data); + sf.options_mut().execution.target_partitions = 11; + info!( + "Set memory limit to {} bytes", + lim.to_formatted_string(&Locale::en) + ); + + let fmm = FairSpillPool::new(lim); + let rt = RuntimeConfig::new().with_memory_pool(Arc::new(fmm)); let ctx = SessionContext::new_with_config(sf); + // let ctx = SessionContext::new_with_config_rt(sf, Arc::new(RuntimeEnv::new(rt)?)); // Register some object store from first input file and output file let store = register_store(store_factory, input_paths, output_path, &ctx)?; diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index 40b93c8f58..c914f762a4 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -16,7 +16,7 @@ * limitations under the License. */ use crate::aws_s3::ObjectStoreFactory; -use aws_config::BehaviorVersion; +use aws_config::{BehaviorVersion, Region}; use aws_credential_types::provider::ProvideCredentials; use color_eyre::eyre::{eyre, Result}; @@ -159,17 +159,19 @@ pub async fn merge_sorted_files(input_data: &CompactionInput<'_>) -> Result Date: Thu, 2 May 2024 13:36:22 +0100 Subject: [PATCH 063/129] Remove memory management code --- rust/Cargo.lock | 57 +++++++++++++++---------------- rust/compaction/Cargo.toml | 1 - rust/compaction/src/datafusion.rs | 16 +-------- rust/compaction/src/details.rs | 5 --- 4 files changed, 29 insertions(+), 50 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 22b2cd9446..10adac0995 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -212,7 +212,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64 0.22.0", + "base64 0.22.1", "chrono", "comfy-table", "half", @@ -422,9 +422,9 @@ checksum = "f1fdabc7756949593fe60f30ec81974b613357de856987752631dea1e3394c80" [[package]] name = "aws-config" -version = "1.2.1" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2a4707646259764ab59fd9a50e9de2e92c637b28b36285d6f6fa030e915fbd9" +checksum = "baaa0be6ee7d90b775ae6ccb6d2ba182b91219ec2001f92338773a094246af1d" dependencies = [ "aws-credential-types", "aws-runtime", @@ -488,9 +488,9 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.21.0" +version = "1.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d70fb493f4183f5102d8a8d0cc9b57aec29a762f55c0e7bf527e0f7177bb408" +checksum = "ca3d6c4cba4e009391b72b0fcf12aff04ea3c9c3aa2ecaafa330326a8bd7e601" dependencies = [ "aws-credential-types", "aws-runtime", @@ -510,9 +510,9 @@ dependencies = [ [[package]] name = "aws-sdk-ssooidc" -version = "1.21.0" +version = "1.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de3f37549b3e38b7ea5efd419d4d7add6ea1e55223506eb0b4fef9d25e7cc90d" +checksum = "73400dc239d14f63d932f4ca7b55af5e9ef1f857f7d70655249ccc287adb2570" dependencies = [ "aws-credential-types", "aws-runtime", @@ -532,9 +532,9 @@ dependencies = [ [[package]] name = "aws-sdk-sts" -version = "1.21.0" +version = "1.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b2ff219a5d4b795cd33251c19dbe9c4b401f2b2cbe513e07c76ada644eaf34e" +checksum = "10f8858308af76fba3e5ffcf1bb56af5471574d2bdfaf0159470c25bc2f760e5" dependencies = [ "aws-credential-types", "aws-runtime", @@ -628,9 +628,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.3.1" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44e7945379821074549168917e89e60630647e186a69243248f08c6d168b975a" +checksum = "1cf64e73ef8d4dac6c933230d56d136b75b252edcf82ed36e37d603090cd7348" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -654,9 +654,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.4.0" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cc56a5c96ec741de6c5e6bf1ce6948be969d6506dfa9c39cffc284e31e4979b" +checksum = "8c19fdae6e3d5ac9cd01f2d6e6c359c5f5a3e028c2d148a8f5b90bf3399a18a7" dependencies = [ "aws-smithy-async", "aws-smithy-types", @@ -739,9 +739,9 @@ checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" [[package]] name = "base64" -version = "0.22.0" +version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9475866fec1451be56a3c2400fd081ff546538961565ccb5b7142cbd22bc7a51" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" [[package]] name = "base64-simd" @@ -868,9 +868,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.0.95" +version = "1.0.96" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d32a725bc159af97c3e629873bb9f88fb8cf8a4867175f76dc987815ea07c83b" +checksum = "065a29261d53ba54260972629f9ca6bffa69bac13cd1fed61420f7fa68b9f8bd" dependencies = [ "jobserver", "libc", @@ -1032,7 +1032,6 @@ dependencies = [ "libc", "log", "num-format", - "num_cpus", "object_store", "rust_sketch", "tokio", @@ -1338,7 +1337,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dd221792c666eac174ecc09e606312844772acc12cbec61a420c2fca1ee70959" dependencies = [ "arrow", - "base64 0.22.0", + "base64 0.22.1", "blake2", "blake3", "chrono", @@ -1407,7 +1406,7 @@ dependencies = [ "arrow-ord", "arrow-schema", "arrow-string", - "base64 0.22.0", + "base64 0.22.1", "blake2", "blake3", "chrono", @@ -1590,9 +1589,9 @@ dependencies = [ [[package]] name = "flate2" -version = "1.0.29" +version = "1.0.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4556222738635b7a3417ae6130d8f52201e45a0c4d1a907f0826383adb5f85e7" +checksum = "5f54427cfd1c7829e2a139fcefea601bf088ebca651d2bf53ebc600eac295dae" dependencies = [ "crc32fast", "miniz_oxide", @@ -2146,9 +2145,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.153" +version = "0.2.154" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" +checksum = "ae743338b92ff9146ce83992f766a31066a91a8c84a45e0e9f21e7cf6de6d346" [[package]] name = "libgit2-sys" @@ -2565,7 +2564,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.22.0", + "base64 0.22.1", "brotli", "bytes", "chrono", @@ -3016,18 +3015,18 @@ checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" [[package]] name = "serde" -version = "1.0.199" +version = "1.0.200" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c9f6e76df036c77cd94996771fb40db98187f096dd0b9af39c6c6e452ba966a" +checksum = "ddc6f9cc94d67c0e21aaf7eda3a010fd3af78ebf6e096aa6e2e13c79749cce4f" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.199" +version = "1.0.200" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11bd257a6541e141e42ca6d24ae26f7714887b47e89aa739099104c7e4d3b7fc" +checksum = "856f046b9400cee3c8c94ed572ecdb752444c24528c035cd35882aad6f492bcb" dependencies = [ "proc-macro2", "quote", diff --git a/rust/compaction/Cargo.toml b/rust/compaction/Cargo.toml index 2371818ad9..f0a7392895 100644 --- a/rust/compaction/Cargo.toml +++ b/rust/compaction/Cargo.toml @@ -53,4 +53,3 @@ tokio-test = { version = "0.4.2" } # Doc tests env_logger = { version = "0.11.3" } chrono = { version = "0.4.26" } # Log helper color-eyre = { version = "0.6.3" } # Error handling -num_cpus = { version = "1.16.0" } # Read CPU count diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index ad7faf4597..043ef17286 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -61,22 +61,10 @@ pub async fn compact( info!("DataFusion compaction of {input_paths:?}"); info!("Compaction region {:?}", input_data.region); - let lim = std::env::var("LIM") - .unwrap_or("104857600".into()) - .parse::() - .unwrap(); - let mut sf = create_session_cfg(input_data); - sf.options_mut().execution.target_partitions = 11; - info!( - "Set memory limit to {} bytes", - lim.to_formatted_string(&Locale::en) - ); + sf.options_mut().execution.target_partitions = 22; - let fmm = FairSpillPool::new(lim); - let rt = RuntimeConfig::new().with_memory_pool(Arc::new(fmm)); let ctx = SessionContext::new_with_config(sf); - // let ctx = SessionContext::new_with_config_rt(sf, Arc::new(RuntimeEnv::new(rt)?)); // Register some object store from first input file and output file let store = register_store(store_factory, input_paths, output_path, &ctx)?; @@ -118,8 +106,6 @@ pub async fn compact( .chain(col_names.iter().skip(1).map(col)) // 1st column is the sketch function call .collect::>(); - let _no_sketches = col_names.iter().map(col).collect::>(); - // Build compaction query frame = frame.sort(sort_order)?.select(col_names_expr)?; diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index c914f762a4..8647969ba5 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -130,11 +130,6 @@ pub struct CompactionResult { /// There must be at least one input file. /// pub async fn merge_sorted_files(input_data: &CompactionInput<'_>) -> Result { - info!( - "Environment has {} CPUs available {} cores", - num_cpus::get_physical(), - num_cpus::get() - ); // Read the schema from the first file if input_data.input_files.is_empty() { Err(eyre!("No input paths supplied")) From 9974ce77295d7cd88a9c34168a852025c2553073 Mon Sep 17 00:00:00 2001 From: m09526 Date: Thu, 2 May 2024 13:57:58 +0100 Subject: [PATCH 064/129] Fix --- rust/compaction/src/details.rs | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/rust/compaction/src/details.rs b/rust/compaction/src/details.rs index cdf341d270..7216648802 100644 --- a/rust/compaction/src/details.rs +++ b/rust/compaction/src/details.rs @@ -16,7 +16,7 @@ * limitations under the License. */ use crate::aws_s3::ObjectStoreFactory; -use aws_config::{BehaviorVersion, Region}; +use aws_config::BehaviorVersion; use aws_credential_types::provider::ProvideCredentials; use color_eyre::eyre::{eyre, Result}; @@ -153,19 +153,17 @@ pub async fn merge_sorted_files(input_data: &CompactionInput<'_>) -> Result Date: Fri, 3 May 2024 08:39:59 +0000 Subject: [PATCH 065/129] Refactor --- .../compaction/job/StateStoreUpdate.java | 42 +++++++++++++++---- ...OutWaitingForFileAssignmentsException.java | 2 +- .../job/execution/StandardCompactor.java | 4 -- ...CompactSortedFilesRetryStateStoreTest.java | 8 ++-- 4 files changed, 41 insertions(+), 15 deletions(-) rename java/compaction/{compaction-job-execution/src/main/java/sleeper/compaction/job/execution => compaction-core/src/main/java/sleeper/compaction/job}/TimedOutWaitingForFileAssignmentsException.java (95%) diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java index 311b3bc030..8c9b626388 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java @@ -21,8 +21,14 @@ import sleeper.core.statestore.FileReference; import sleeper.core.statestore.StateStore; import sleeper.core.statestore.StateStoreException; +import sleeper.core.statestore.exception.FileReferenceNotAssignedToJobException; +import sleeper.core.util.ExponentialBackoffWithJitter; +import sleeper.core.util.ExponentialBackoffWithJitter.WaitRange; public class StateStoreUpdate { + public static final int JOB_ASSIGNMENT_WAIT_ATTEMPTS = 10; + public static final WaitRange JOB_ASSIGNMENT_WAIT_RANGE = WaitRange.firstAndMaxWaitCeilingSecs(2, 60); + private StateStoreUpdate() { } @@ -31,7 +37,18 @@ private StateStoreUpdate() { public static void updateStateStoreSuccess( CompactionJob job, long recordsWritten, - StateStore stateStore) throws StateStoreException { + StateStore stateStore) throws StateStoreException, InterruptedException { + updateStateStoreSuccess(job, recordsWritten, stateStore, + JOB_ASSIGNMENT_WAIT_ATTEMPTS, + new ExponentialBackoffWithJitter(JOB_ASSIGNMENT_WAIT_RANGE)); + } + + public static void updateStateStoreSuccess( + CompactionJob job, + long recordsWritten, + StateStore stateStore, + int jobAssignmentWaitAttempts, + ExponentialBackoffWithJitter jobAssignmentWaitBackoff) throws StateStoreException, InterruptedException { FileReference fileReference = FileReference.builder() .filename(job.getOutputFile()) .partitionId(job.getPartitionId()) @@ -39,12 +56,23 @@ public static void updateStateStoreSuccess( .countApproximate(false) .onlyContainsDataForThisPartition(true) .build(); - try { - stateStore.atomicallyReplaceFileReferencesWithNewOne(job.getId(), job.getPartitionId(), job.getInputFiles(), fileReference); - LOGGER.debug("Updated file references in state store"); - } catch (StateStoreException e) { - LOGGER.error("Exception updating StateStore (moving input files to ready for GC and creating new active file): {}", e.getMessage()); - throw e; + + // Compaction jobs are sent for execution before updating the state store to assign the input files to the job. + // Sometimes the compaction can finish before the job assignment is finished. We wait for the job assignment + // rather than immediately failing the job run. + FileReferenceNotAssignedToJobException failure = null; + for (int attempts = 0; attempts < jobAssignmentWaitAttempts; attempts++) { + jobAssignmentWaitBackoff.waitBeforeAttempt(attempts); + try { + stateStore.atomicallyReplaceFileReferencesWithNewOne(job.getId(), job.getPartitionId(), job.getInputFiles(), fileReference); + LOGGER.debug("Updated file references in state store"); + return; + } catch (FileReferenceNotAssignedToJobException e) { + LOGGER.warn("Job not yet assigned to input files on attempt {} of {}: {}", + attempts + 1, jobAssignmentWaitAttempts, e.getMessage()); + failure = e; + } } + throw new TimedOutWaitingForFileAssignmentsException(failure); } } diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/TimedOutWaitingForFileAssignmentsException.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java similarity index 95% rename from java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/TimedOutWaitingForFileAssignmentsException.java rename to java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java index 030af4a104..f0a1bc5cec 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/TimedOutWaitingForFileAssignmentsException.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.execution; +package sleeper.compaction.job; public class TimedOutWaitingForFileAssignmentsException extends RuntimeException { diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java index d6dd92c77f..a8e3285195 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java @@ -42,7 +42,6 @@ import sleeper.core.schema.Schema; import sleeper.core.statestore.StateStore; import sleeper.core.statestore.StateStoreException; -import sleeper.core.util.ExponentialBackoffWithJitter.WaitRange; import sleeper.io.parquet.record.ParquetReaderIterator; import sleeper.io.parquet.record.ParquetRecordReader; import sleeper.io.parquet.record.ParquetRecordWriterFactory; @@ -65,9 +64,6 @@ * Executes a compaction job. Compacts N input files into a single output file. */ public class StandardCompactor implements CompactionRunner { - public static final int JOB_ASSIGNMENT_WAIT_ATTEMPTS = 10; - public static final WaitRange JOB_ASSIGNMENT_WAIT_RANGE = WaitRange.firstAndMaxWaitCeilingSecs(2, 60); - private final InstanceProperties instanceProperties; private final TablePropertiesProvider tablePropertiesProvider; private final ObjectFactory objectFactory; diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java index 31e1566d08..ec8ff627c8 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java @@ -20,6 +20,8 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionJobFactory; +import sleeper.compaction.job.StateStoreUpdate; +import sleeper.compaction.job.TimedOutWaitingForFileAssignmentsException; import sleeper.configuration.properties.instance.InstanceProperties; import sleeper.configuration.properties.table.TableProperties; import sleeper.core.partition.PartitionTree; @@ -130,8 +132,8 @@ void shouldFailWithNoRetriesWhenFileDoesNotExist() throws Exception { } private void updateStateStoreSuccess(CompactionJob job, long recordsWritten, DoubleSupplier randomJitter) throws Exception { - CompactSortedFiles.updateStateStoreSuccess(job, 123, stateStore, - CompactSortedFiles.JOB_ASSIGNMENT_WAIT_ATTEMPTS, backoff(randomJitter)); + StateStoreUpdate.updateStateStoreSuccess(job, 123, stateStore, + StateStoreUpdate.JOB_ASSIGNMENT_WAIT_ATTEMPTS, backoff(randomJitter)); } private void actionOnWait(WaitAction action) throws Exception { @@ -152,7 +154,7 @@ private CompactionJob createCompactionJobForOneFile(FileReference file) { private ExponentialBackoffWithJitter backoff(DoubleSupplier randomJitter) { return new ExponentialBackoffWithJitter( - CompactSortedFiles.JOB_ASSIGNMENT_WAIT_RANGE, + StateStoreUpdate.JOB_ASSIGNMENT_WAIT_RANGE, randomJitter, waiter); } From 3d7f3b36bd04aa3eb97e8e9010b3fd23ce849837 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 3 May 2024 10:16:40 +0000 Subject: [PATCH 066/129] Minor improvements to DF --- rust/compaction/src/datafusion.rs | 23 +++++++++----- rust/compaction/src/datafusion/sketch.rs | 39 ++++++++++++++++++++++++ rust/compaction/src/datafusion/udf.rs | 1 - rust/compaction/src/lib.rs | 4 +-- 4 files changed, 56 insertions(+), 11 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 5c9406ebf8..1f86885475 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -28,9 +28,7 @@ use datafusion::{ dataframe::DataFrameWriteOptions, error::DataFusionError, execution::{ - config::SessionConfig, - context::SessionContext, - options::ParquetReadOptions, + config::SessionConfig, context::SessionContext, options::ParquetReadOptions, FunctionRegistry, }, logical_expr::ScalarUDF, @@ -56,10 +54,14 @@ pub async fn compact( input_paths: &[Url], output_path: &Url, ) -> Result { - info!("DataFusion compaction of {input_paths:?}"); + info!( + "DataFusion compaction of files {:?}", + input_paths.iter().map(Url::as_str).collect::>() + ); + info!("DataFusion output file {}", output_path.as_str()); info!("Compaction partition region {:?}", input_data.region); - let sf = create_session_cfg(input_data); + let sf = create_session_cfg(input_data, input_paths); let ctx = SessionContext::new_with_config(sf); // Register some object store from first input file and output file @@ -139,10 +141,14 @@ pub async fn compact( let binding = sketch_func.inner(); let inner_function: Option<&SketchUDF> = binding.as_any().downcast_ref(); if let Some(func) = inner_function { + let first_sketch = &func.get_sketch()[0]; info!( - "Made {} calls to sketch UDF and processed {} total rows.", + "Made {} calls to sketch UDF and processed {} total rows. Quantile sketch column 0 retained {} out of {} values (K value = {}).", func.get_invoke_count().to_formatted_string(&Locale::en), func.get_row_count().to_formatted_string(&Locale::en), + first_sketch.get_num_retained(), + first_sketch.get_n(), + first_sketch.get_k() ); rows_written = func.get_row_count(); @@ -292,11 +298,12 @@ fn get_parquet_writer_version(version: &str) -> String { /// /// This sets as many parameters as possible from the given input data. /// -fn create_session_cfg(input_data: &CompactionInput) -> SessionConfig { +fn create_session_cfg(input_data: &CompactionInput, input_paths: &[T]) -> SessionConfig { let mut sf = SessionConfig::new(); // In order to avoid a costly "Sort" stage in the physical plan, we must make // sure the target partitions as at least as big as number of input files. - sf.options_mut().execution.target_partitions = std::cmp::max(sf.options().execution.target_partitions, input_data.input_files.len()); + sf.options_mut().execution.target_partitions = + std::cmp::max(sf.options().execution.target_partitions, input_paths.len()); sf.options_mut().execution.parquet.max_row_group_size = input_data.max_row_group_size; sf.options_mut().execution.parquet.data_pagesize_limit = input_data.max_page_size; sf.options_mut().execution.parquet.compression = Some(get_compression(&input_data.compression)); diff --git a/rust/compaction/src/datafusion/sketch.rs b/rust/compaction/src/datafusion/sketch.rs index 89debce5af..09a913c69b 100644 --- a/rust/compaction/src/datafusion/sketch.rs +++ b/rust/compaction/src/datafusion/sketch.rs @@ -182,6 +182,45 @@ impl DataSketchVariant { } } + /// Gets the 'k' parameter of the quantile sketch. + /// + /// Please see Apache data sketch C++ documentation for full explanation. + pub fn get_k(&self) -> u16 { + match self { + DataSketchVariant::I32(s) => s.get_k(), + DataSketchVariant::I64(s) => s.get_k(), + DataSketchVariant::Str(_, s) => s.get_k(), + DataSketchVariant::Bytes(_, s) => s.get_k(), + } + } + + /// Gets the total number of items in this sketch. + /// + /// Note that as this sketches are approximate, only a fraction of this amount + /// is retained by the sketch. Please see [`get_num_retained`]. + /// + /// Please see Apache data sketch C++ documentation for full explanation. + pub fn get_n(&self) -> u64 { + match self { + DataSketchVariant::I32(s) => s.get_n(), + DataSketchVariant::I64(s) => s.get_n(), + DataSketchVariant::Str(_, s) => s.get_n(), + DataSketchVariant::Bytes(_, s) => s.get_n(), + } + } + + /// Gets the number of individual items retained by the sketch. + /// + /// Please see Apache data sketch C++ documentation for full explanation. + pub fn get_num_retained(&self) -> u32 { + match self { + DataSketchVariant::I32(s) => s.get_num_retained(), + DataSketchVariant::I64(s) => s.get_num_retained(), + DataSketchVariant::Str(_, s) => s.get_num_retained(), + DataSketchVariant::Bytes(_, s) => s.get_num_retained(), + } + } + /// Get the minimum item from this sketch. /// /// # Errors diff --git a/rust/compaction/src/datafusion/udf.rs b/rust/compaction/src/datafusion/udf.rs index fb275e80fb..862cd3ce9f 100644 --- a/rust/compaction/src/datafusion/udf.rs +++ b/rust/compaction/src/datafusion/udf.rs @@ -171,7 +171,6 @@ impl ScalarUDFImpl for SketchUDF { ), _ => return internal_err!("Row type {} not supported for Sleeper row key field", array.data_type()), } - *self.row_count.lock().unwrap() += array.len(); } ColumnarValue::Scalar( diff --git a/rust/compaction/src/lib.rs b/rust/compaction/src/lib.rs index 279b33432d..4ad22fae76 100644 --- a/rust/compaction/src/lib.rs +++ b/rust/compaction/src/lib.rs @@ -215,7 +215,7 @@ pub extern "C" fn allocate_result() -> *const FFICompactionResult { rows_read: 0, rows_written: 0, })); - info!("Compaction result allocated @ {:p}", p); + info!("Compaction result allocated at address {:p}", p); p } @@ -443,7 +443,7 @@ fn unpack_variant_array<'a>( pub extern "C" fn free_result(ob: *mut FFICompactionResult) { maybe_cfg_log(); if !ob.is_null() { - info!("Compaction result destructed at {:p}", ob); + info!("Compaction result at address {:p} destructed", ob); let _ = unsafe { Box::from_raw(ob) }; } } From b1565eaf2ad133bd26a74a369525ad7312b110e8 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 3 May 2024 12:06:00 +0000 Subject: [PATCH 067/129] Deadlock removed --- rust/compaction/src/datafusion.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 1f86885475..5c9d6199b2 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -141,8 +141,10 @@ pub async fn compact( let binding = sketch_func.inner(); let inner_function: Option<&SketchUDF> = binding.as_any().downcast_ref(); if let Some(func) = inner_function { - let first_sketch = &func.get_sketch()[0]; - info!( + { + // Limit scope of MutexGuard + let first_sketch = &func.get_sketch()[0]; + info!( "Made {} calls to sketch UDF and processed {} total rows. Quantile sketch column 0 retained {} out of {} values (K value = {}).", func.get_invoke_count().to_formatted_string(&Locale::en), func.get_row_count().to_formatted_string(&Locale::en), @@ -150,6 +152,7 @@ pub async fn compact( first_sketch.get_n(), first_sketch.get_k() ); + } rows_written = func.get_row_count(); From c3c72e661ddefccc8b6132f5b5e6119f8a22e18e Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 3 May 2024 12:07:06 +0000 Subject: [PATCH 068/129] Typo --- rust/compaction/src/datafusion.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 5c9d6199b2..8efead8261 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -145,7 +145,7 @@ pub async fn compact( // Limit scope of MutexGuard let first_sketch = &func.get_sketch()[0]; info!( - "Made {} calls to sketch UDF and processed {} total rows. Quantile sketch column 0 retained {} out of {} values (K value = {}).", + "Made {} calls to sketch UDF and processed {} rows. Quantile sketch column 0 retained {} out of {} values (K value = {}).", func.get_invoke_count().to_formatted_string(&Locale::en), func.get_row_count().to_formatted_string(&Locale::en), first_sketch.get_num_retained(), From 19cc6c5e360be522872bce208d32d4ad271e91b5 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 3 May 2024 12:08:25 +0000 Subject: [PATCH 069/129] Formatting --- rust/compaction/src/datafusion.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/rust/compaction/src/datafusion.rs b/rust/compaction/src/datafusion.rs index 8efead8261..241f58361c 100644 --- a/rust/compaction/src/datafusion.rs +++ b/rust/compaction/src/datafusion.rs @@ -148,9 +148,9 @@ pub async fn compact( "Made {} calls to sketch UDF and processed {} rows. Quantile sketch column 0 retained {} out of {} values (K value = {}).", func.get_invoke_count().to_formatted_string(&Locale::en), func.get_row_count().to_formatted_string(&Locale::en), - first_sketch.get_num_retained(), - first_sketch.get_n(), - first_sketch.get_k() + first_sketch.get_num_retained().to_formatted_string(&Locale::en), + first_sketch.get_n().to_formatted_string(&Locale::en), + first_sketch.get_k().to_formatted_string(&Locale::en) ); } From 3996446a49bb15abaf9c5cbf19ee6900a3b9295a Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 3 May 2024 14:16:49 +0000 Subject: [PATCH 070/129] Uppercase selection --- .../java/sleeper/compaction/job/execution/DefaultSelector.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java index 39be52c8f2..2350277394 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/DefaultSelector.java @@ -52,7 +52,7 @@ public DefaultSelector( public CompactionRunner chooseCompactor(CompactionJob job) { TableProperties tableProperties = tablePropertiesProvider .getById(job.getTableId()); - String method = tableProperties.get(TableProperty.COMPACTION_METHOD); + String method = tableProperties.get(TableProperty.COMPACTION_METHOD).toUpperCase(); // Convert to enum value and default to Java CompactionMethod desired; From f1a511efd04efdb0eab432f8385ad1291cf9515a Mon Sep 17 00:00:00 2001 From: m09526 Date: Wed, 8 May 2024 12:22:59 +0000 Subject: [PATCH 071/129] Moved --- ...OutWaitingForFileAssignmentsException.java | 24 ------------------- 1 file changed, 24 deletions(-) delete mode 100644 java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java deleted file mode 100644 index f0a1bc5cec..0000000000 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java +++ /dev/null @@ -1,24 +0,0 @@ -/* - * 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.compaction.job; - -public class TimedOutWaitingForFileAssignmentsException extends RuntimeException { - - public TimedOutWaitingForFileAssignmentsException(Throwable cause) { - super("Too many retries waiting for input files to be assigned to job in state store", cause); - } - -} From 18c8b66430263ac6f07c6846a4edfb2cf4c0e4d3 Mon Sep 17 00:00:00 2001 From: m09526 Date: Wed, 8 May 2024 12:33:59 +0000 Subject: [PATCH 072/129] Checkstyle --- .../main/java/sleeper/compaction/job/StateStoreUpdate.java | 1 + .../TimedOutWaitingForFileAssignmentsException.java | 4 ---- .../sleeper/compaction/job/execution/StandardCompactor.java | 1 - 3 files changed, 1 insertion(+), 5 deletions(-) diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java index 8c9b626388..a20b997c89 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/StateStoreUpdate.java @@ -18,6 +18,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import sleeper.compaction.job.completion.TimedOutWaitingForFileAssignmentsException; import sleeper.core.statestore.FileReference; import sleeper.core.statestore.StateStore; import sleeper.core.statestore.StateStoreException; diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java index 3a5decf577..dec7ef5983 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java @@ -13,11 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -<<<<<<<< HEAD:java/compaction/compaction-core/src/main/java/sleeper/compaction/job/TimedOutWaitingForFileAssignmentsException.java -package sleeper.compaction.job; -======== package sleeper.compaction.job.completion; ->>>>>>>> develop:java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java public class TimedOutWaitingForFileAssignmentsException extends RuntimeException { diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java index 74ebf4fbf0..a8e3285195 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java @@ -27,7 +27,6 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionRunner; import sleeper.compaction.job.StateStoreUpdate; -import sleeper.compaction.job.completion.TimedOutWaitingForFileAssignmentsException; import sleeper.configuration.jars.ObjectFactory; import sleeper.configuration.jars.ObjectFactoryException; import sleeper.configuration.properties.instance.InstanceProperties; From a28df4a050525e6d69c4f2711315e292a378e2c9 Mon Sep 17 00:00:00 2001 From: m09526 Date: Fri, 10 May 2024 11:26:53 +0100 Subject: [PATCH 073/129] Merge develop --- .github/config/chunks.yaml | 2 +- .github/workflows/chunk-compaction.yaml | 2 +- code-style/dependency-check-suppressions.xml | 2 +- example/full/instance.properties | 5 + example/full/table.properties | 5 + .../CompactionJobCommitRequest.java} | 12 +- .../CompactionJobCommitRequestSerDe.java} | 14 +- .../job/commit/CompactionJobCommitter.java | 67 +++ .../CompactionJobCommitterUtils.java} | 56 +- ...OutWaitingForFileAssignmentsException.java | 2 +- .../CompactionJobCommitRequestSerDeTest.java} | 14 +- ...seCompactionJobCommitRequest.approved.txt} | 0 .../CompactionJobCommitterTest.java} | 28 +- .../CompactionJobCommitterTestBase.java} | 22 +- .../pom.xml | 2 +- .../lambda/CompactionJobCommitterLambda.java} | 38 +- .../execution/CompactionJobCommitHandler.java | 55 ++ .../job/execution/CompactionTask.java | 16 +- .../execution/ECSCompactionTaskRunner.java | 16 +- .../job/execution/StandardCompactor.java | 14 +- .../CompactSortedFilesEmptyOutputIT.java | 22 - .../job/execution/CompactSortedFilesIT.java | 39 +- .../CompactSortedFilesIteratorIT.java | 15 +- .../CompactSortedFilesLocalStackIT.java | 14 +- ...CompactSortedFilesRetryStateStoreTest.java | 13 +- .../execution/CompactionTaskCommitTest.java | 313 ++++++++++ .../CompactionTaskTerminateTest.java | 218 +++++++ .../job/execution/CompactionTaskTest.java | 540 +----------------- .../job/execution/CompactionTaskTestBase.java | 312 ++++++++++ .../ECSCompactionTaskRunnerLocalStackIT.java | 12 +- java/compaction/pom.xml | 2 +- .../properties/instance/DefaultProperty.java | 6 + .../properties/table/TableProperty.java | 7 + java/pom.xml | 4 +- .../testutil/drivers/InMemoryCompaction.java | 1 - scripts/templates/instanceproperties.template | 5 + 36 files changed, 1131 insertions(+), 764 deletions(-) rename java/compaction/compaction-core/src/main/java/sleeper/compaction/job/{completion/CompactionJobCompletionRequest.java => commit/CompactionJobCommitRequest.java} (84%) rename java/compaction/compaction-core/src/main/java/sleeper/compaction/job/{completion/CompactionJobCompletionRequestSerDe.java => commit/CompactionJobCommitRequestSerDe.java} (85%) create mode 100644 java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitter.java rename java/compaction/compaction-core/src/main/java/sleeper/compaction/job/{completion/CompactionJobCompletion.java => commit/CompactionJobCommitterUtils.java} (53%) rename java/compaction/compaction-core/src/main/java/sleeper/compaction/job/{completion => commit}/TimedOutWaitingForFileAssignmentsException.java (95%) rename java/compaction/compaction-core/src/test/java/sleeper/compaction/job/{completion/CompactionJobCompletionRequestSerDeTest.java => commit/CompactionJobCommitRequestSerDeTest.java} (76%) rename java/compaction/compaction-core/src/test/java/sleeper/compaction/job/{completion/CompactionJobCompletionRequestSerDeTest.shouldSerialiseCompactionJobCompletion.approved.txt => commit/CompactionJobCommitRequestSerDeTest.shouldSerialiseCompactionJobCommitRequest.approved.txt} (100%) rename java/compaction/compaction-core/src/test/java/sleeper/compaction/job/{completion/CompactionJobCompletionTest.java => commit/CompactionJobCommitterTest.java} (86%) rename java/compaction/compaction-core/src/test/java/sleeper/compaction/job/{completion/CompactionJobCompletionTestBase.java => commit/CompactionJobCommitterTestBase.java} (89%) rename java/compaction/{compaction-job-completion-lambda => compaction-job-committer-lambda}/pom.xml (98%) rename java/compaction/{compaction-job-completion-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCompletionLambda.java => compaction-job-committer-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCommitterLambda.java} (75%) create mode 100644 java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionJobCommitHandler.java create mode 100644 java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskCommitTest.java create mode 100644 java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTerminateTest.java create mode 100644 java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTestBase.java diff --git a/.github/config/chunks.yaml b/.github/config/chunks.yaml index 2c816ced83..d098c4d203 100644 --- a/.github/config/chunks.yaml +++ b/.github/config/chunks.yaml @@ -22,7 +22,7 @@ chunks: compaction: name: Compaction workflow: chunk-compaction.yaml - modules: [ compaction/compaction-job-execution, compaction/compaction-task-creation, compaction/compaction-job-creation, compaction/compaction-job-creation-lambda, compaction/compaction-job-completion-lambda, compaction/compaction-rust, compaction/compaction-status-store, compaction/compaction-core, splitter/splitter-core, splitter/splitter-lambda ] + modules: [ compaction/compaction-job-execution, compaction/compaction-task-creation, compaction/compaction-job-creation, compaction/compaction-job-creation-lambda, compaction/compaction-job-committer-lambda, compaction/compaction-status-store, compaction/compaction-core, splitter/splitter-core, splitter/splitter-lambda, compaction/compaction-rust, ] data: name: Data workflow: chunk-data.yaml diff --git a/.github/workflows/chunk-compaction.yaml b/.github/workflows/chunk-compaction.yaml index 436be2bfac..c8463bcb04 100644 --- a/.github/workflows/chunk-compaction.yaml +++ b/.github/workflows/chunk-compaction.yaml @@ -15,7 +15,7 @@ on: - 'java/compaction/compaction-rust/**' - 'java/compaction/compaction-job-creation/**' - 'java/compaction/compaction-job-creation-lambda/**' - - 'java/compaction/compaction-job-completion-lambda/**' + - 'java/compaction/compaction-job-committer-lambda/**' - 'java/compaction/compaction-status-store/**' - 'java/compaction/compaction-core/**' - 'java/splitter/splitter-core/**' diff --git a/code-style/dependency-check-suppressions.xml b/code-style/dependency-check-suppressions.xml index 23eafb7218..4236d12060 100644 --- a/code-style/dependency-check-suppressions.xml +++ b/code-style/dependency-check-suppressions.xml @@ -101,7 +101,7 @@ doesn't have a fixed version. ]]> ^pkg:maven/org\.bouncycastle/bcprov\-jdk15on@.*$ - CVE-2023-33201|CVE-2024-29857|CVE-2024-30172|CVE-2024-30171 + CVE-2023-33201|CVE-2024-29857|CVE-2024-30172|CVE-2024-30171|CVE-2024-34447 = ratio * # s_n. diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletionRequest.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitRequest.java similarity index 84% rename from java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletionRequest.java rename to java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitRequest.java index 5bb43bedf7..1e83fd8ec2 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletionRequest.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitRequest.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.completion; +package sleeper.compaction.job.commit; import sleeper.compaction.job.CompactionJob; import sleeper.core.record.process.RecordsProcessed; @@ -22,7 +22,7 @@ import java.time.Instant; import java.util.Objects; -public class CompactionJobCompletionRequest { +public class CompactionJobCommitRequest { private final CompactionJob job; private final String taskId; @@ -31,7 +31,7 @@ public class CompactionJobCompletionRequest { private final long recordsRead; private final long recordsWritten; - public CompactionJobCompletionRequest( + public CompactionJobCommitRequest( CompactionJob job, String taskId, RecordsProcessedSummary recordsProcessed) { this.job = job; this.taskId = taskId; @@ -75,17 +75,17 @@ public boolean equals(Object obj) { if (this == obj) { return true; } - if (!(obj instanceof CompactionJobCompletionRequest)) { + if (!(obj instanceof CompactionJobCommitRequest)) { return false; } - CompactionJobCompletionRequest other = (CompactionJobCompletionRequest) obj; + CompactionJobCommitRequest other = (CompactionJobCommitRequest) obj; return Objects.equals(job, other.job) && Objects.equals(taskId, other.taskId) && Objects.equals(startTime, other.startTime) && Objects.equals(finishTime, other.finishTime) && recordsRead == other.recordsRead && recordsWritten == other.recordsWritten; } @Override public String toString() { - return "CompactionJobCompletionRequest{job=" + job + ", taskId=" + taskId + ", startTime=" + startTime + ", finishTime=" + finishTime + ", recordsRead=" + recordsRead + ", recordsWritten=" + return "CompactionJobCommitRequest{job=" + job + ", taskId=" + taskId + ", startTime=" + startTime + ", finishTime=" + finishTime + ", recordsRead=" + recordsRead + ", recordsWritten=" + recordsWritten + "}"; } diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDe.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDe.java similarity index 85% rename from java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDe.java rename to java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDe.java index 6d67e1ab0d..9ecacf4736 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDe.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDe.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.completion; +package sleeper.compaction.job.commit; import com.google.gson.Gson; import com.google.gson.GsonBuilder; @@ -33,12 +33,12 @@ import java.io.UncheckedIOException; import java.lang.reflect.Type; -public class CompactionJobCompletionRequestSerDe { +public class CompactionJobCommitRequestSerDe { private final Gson gson; private final Gson gsonPrettyPrint; - public CompactionJobCompletionRequestSerDe() { + public CompactionJobCommitRequestSerDe() { GsonBuilder builder = GsonConfig.standardBuilder() .registerTypeAdapter(CompactionJob.class, new CompactionJobJsonSerDe()) .serializeNulls(); @@ -46,16 +46,16 @@ public CompactionJobCompletionRequestSerDe() { gsonPrettyPrint = builder.setPrettyPrinting().create(); } - public String toJson(CompactionJobCompletionRequest jobRun) { + public String toJson(CompactionJobCommitRequest jobRun) { return gson.toJson(jobRun); } - public String toJsonPrettyPrint(CompactionJobCompletionRequest jobRun) { + public String toJsonPrettyPrint(CompactionJobCommitRequest jobRun) { return gsonPrettyPrint.toJson(jobRun); } - public CompactionJobCompletionRequest fromJson(String json) { - return gson.fromJson(json, CompactionJobCompletionRequest.class); + public CompactionJobCommitRequest fromJson(String json) { + return gson.fromJson(json, CompactionJobCommitRequest.class); } /** diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitter.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitter.java new file mode 100644 index 0000000000..e56b7c289d --- /dev/null +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitter.java @@ -0,0 +1,67 @@ +/* + * 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.compaction.job.commit; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import sleeper.compaction.job.CompactionJob; +import sleeper.compaction.job.CompactionJobStatusStore; +import sleeper.core.statestore.StateStore; +import sleeper.core.statestore.StateStoreException; +import sleeper.core.util.ExponentialBackoffWithJitter; +import sleeper.core.util.ExponentialBackoffWithJitter.WaitRange; + +import static sleeper.compaction.job.commit.CompactionJobCommitterUtils.updateStateStoreSuccess; + +public class CompactionJobCommitter { + public static final Logger LOGGER = LoggerFactory.getLogger(CompactionJobCommitter.class); + + public static final int JOB_ASSIGNMENT_WAIT_ATTEMPTS = 10; + public static final WaitRange JOB_ASSIGNMENT_WAIT_RANGE = WaitRange.firstAndMaxWaitCeilingSecs(2, 60); + + private final CompactionJobStatusStore statusStore; + private final GetStateStore stateStoreProvider; + private final int jobAssignmentWaitAttempts; + private final ExponentialBackoffWithJitter jobAssignmentWaitBackoff; + + public CompactionJobCommitter( + CompactionJobStatusStore statusStore, GetStateStore stateStoreProvider) { + this(statusStore, stateStoreProvider, JOB_ASSIGNMENT_WAIT_ATTEMPTS, + new ExponentialBackoffWithJitter(JOB_ASSIGNMENT_WAIT_RANGE)); + } + + public CompactionJobCommitter( + CompactionJobStatusStore statusStore, GetStateStore stateStoreProvider, + int jobAssignmentWaitAttempts, ExponentialBackoffWithJitter jobAssignmentWaitBackoff) { + this.statusStore = statusStore; + this.stateStoreProvider = stateStoreProvider; + this.jobAssignmentWaitAttempts = jobAssignmentWaitAttempts; + this.jobAssignmentWaitBackoff = jobAssignmentWaitBackoff; + } + + public void apply(CompactionJobCommitRequest request) throws StateStoreException, InterruptedException { + CompactionJob job = request.getJob(); + updateStateStoreSuccess(job, request.getRecordsWritten(), stateStoreProvider.getByTableId(job.getTableId()), + jobAssignmentWaitAttempts, jobAssignmentWaitBackoff); + statusStore.jobFinished(job, request.buildRecordsProcessedSummary(), request.getTaskId()); + } + + @FunctionalInterface + public interface GetStateStore { + StateStore getByTableId(String tableId); + } +} diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletion.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitterUtils.java similarity index 53% rename from java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletion.java rename to java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitterUtils.java index 7bd2c2a39a..a69f3b0e40 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/CompactionJobCompletion.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/CompactionJobCommitterUtils.java @@ -13,13 +13,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.completion; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +package sleeper.compaction.job.commit; import sleeper.compaction.job.CompactionJob; -import sleeper.compaction.job.CompactionJobStatusStore; import sleeper.core.statestore.FileReference; import sleeper.core.statestore.StateStore; import sleeper.core.statestore.StateStoreException; @@ -27,40 +23,28 @@ import sleeper.core.util.ExponentialBackoffWithJitter; import sleeper.core.util.ExponentialBackoffWithJitter.WaitRange; -public class CompactionJobCompletion { - public static final Logger LOGGER = LoggerFactory.getLogger(CompactionJobCompletion.class); - +public class CompactionJobCommitterUtils { public static final int JOB_ASSIGNMENT_WAIT_ATTEMPTS = 10; public static final WaitRange JOB_ASSIGNMENT_WAIT_RANGE = WaitRange.firstAndMaxWaitCeilingSecs(2, 60); - private final CompactionJobStatusStore statusStore; - private final GetStateStore stateStoreProvider; - private final int jobAssignmentWaitAttempts; - private final ExponentialBackoffWithJitter jobAssignmentWaitBackoff; - - public CompactionJobCompletion( - CompactionJobStatusStore statusStore, GetStateStore stateStoreProvider) { - this(statusStore, stateStoreProvider, JOB_ASSIGNMENT_WAIT_ATTEMPTS, - new ExponentialBackoffWithJitter(JOB_ASSIGNMENT_WAIT_RANGE)); + private CompactionJobCommitterUtils() { } - public CompactionJobCompletion( - CompactionJobStatusStore statusStore, GetStateStore stateStoreProvider, - int jobAssignmentWaitAttempts, ExponentialBackoffWithJitter jobAssignmentWaitBackoff) { - this.statusStore = statusStore; - this.stateStoreProvider = stateStoreProvider; - this.jobAssignmentWaitAttempts = jobAssignmentWaitAttempts; - this.jobAssignmentWaitBackoff = jobAssignmentWaitBackoff; - } - - public void apply(CompactionJobCompletionRequest request) throws StateStoreException, InterruptedException { - CompactionJob job = request.getJob(); - updateStateStoreSuccess(job, request.getRecordsWritten()); - statusStore.jobFinished(job, request.buildRecordsProcessedSummary(), request.getTaskId()); + public static void updateStateStoreSuccess( + CompactionJob job, + long recordsWritten, + StateStore stateStore) throws StateStoreException, InterruptedException { + updateStateStoreSuccess(job, recordsWritten, stateStore, + JOB_ASSIGNMENT_WAIT_ATTEMPTS, + new ExponentialBackoffWithJitter(JOB_ASSIGNMENT_WAIT_RANGE)); } - private void updateStateStoreSuccess(CompactionJob job, long recordsWritten) throws StateStoreException, InterruptedException { - StateStore stateStore = stateStoreProvider.getByTableId(job.getTableId()); + public static void updateStateStoreSuccess( + CompactionJob job, + long recordsWritten, + StateStore stateStore, + int jobAssignmentWaitAttempts, + ExponentialBackoffWithJitter jobAssignmentWaitBackoff) throws StateStoreException, InterruptedException { FileReference fileReference = FileReference.builder() .filename(job.getOutputFile()) .partitionId(job.getPartitionId()) @@ -77,19 +61,11 @@ private void updateStateStoreSuccess(CompactionJob job, long recordsWritten) thr jobAssignmentWaitBackoff.waitBeforeAttempt(attempts); try { stateStore.atomicallyReplaceFileReferencesWithNewOne(job.getId(), job.getPartitionId(), job.getInputFiles(), fileReference); - LOGGER.info("Atomically replaced {} file references in state store with file reference {}.", job.getInputFiles(), fileReference); return; } catch (FileReferenceNotAssignedToJobException e) { - LOGGER.warn("Job not yet assigned to input files on attempt {} of {}: {}", - attempts + 1, jobAssignmentWaitAttempts, e.getMessage()); failure = e; } } throw new TimedOutWaitingForFileAssignmentsException(failure); } - - @FunctionalInterface - public interface GetStateStore { - StateStore getByTableId(String tableId); - } } diff --git a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/TimedOutWaitingForFileAssignmentsException.java similarity index 95% rename from java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java rename to java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/TimedOutWaitingForFileAssignmentsException.java index dec7ef5983..fa22acf574 100644 --- a/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/completion/TimedOutWaitingForFileAssignmentsException.java +++ b/java/compaction/compaction-core/src/main/java/sleeper/compaction/job/commit/TimedOutWaitingForFileAssignmentsException.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.completion; +package sleeper.compaction.job.commit; public class TimedOutWaitingForFileAssignmentsException extends RuntimeException { diff --git a/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDeTest.java b/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDeTest.java similarity index 76% rename from java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDeTest.java rename to java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDeTest.java index 805a0e740c..54e2f0e28b 100644 --- a/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDeTest.java +++ b/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDeTest.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.completion; +package sleeper.compaction.job.commit; import org.approvaltests.Approvals; import org.junit.jupiter.api.Test; @@ -28,12 +28,12 @@ import static org.assertj.core.api.Assertions.assertThat; -public class CompactionJobCompletionRequestSerDeTest { +public class CompactionJobCommitRequestSerDeTest { - private final CompactionJobCompletionRequestSerDe serDe = new CompactionJobCompletionRequestSerDe(); + private final CompactionJobCommitRequestSerDe serDe = new CompactionJobCommitRequestSerDe(); @Test - void shouldSerialiseCompactionJobCompletion() throws Exception { + void shouldSerialiseCompactionJobCommitRequest() throws Exception { // Given CompactionJob job = CompactionJob.builder() .tableId("test-table") @@ -42,16 +42,16 @@ void shouldSerialiseCompactionJobCompletion() throws Exception { .outputFile("test-output.parquet") .partitionId("test-partition-id") .build(); - CompactionJobCompletionRequest completion = new CompactionJobCompletionRequest(job, "test-task", + CompactionJobCommitRequest commit = new CompactionJobCommitRequest(job, "test-task", new RecordsProcessedSummary( new RecordsProcessed(120, 100), Instant.parse("2024-05-01T10:58:00Z"), Duration.ofMinutes(1))); // When - String json = serDe.toJsonPrettyPrint(completion); + String json = serDe.toJsonPrettyPrint(commit); // Then - assertThat(serDe.fromJson(json)).isEqualTo(completion); + assertThat(serDe.fromJson(json)).isEqualTo(commit); Approvals.verify(json); } } diff --git a/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDeTest.shouldSerialiseCompactionJobCompletion.approved.txt b/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDeTest.shouldSerialiseCompactionJobCommitRequest.approved.txt similarity index 100% rename from java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionRequestSerDeTest.shouldSerialiseCompactionJobCompletion.approved.txt rename to java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitRequestSerDeTest.shouldSerialiseCompactionJobCommitRequest.approved.txt diff --git a/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionTest.java b/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitterTest.java similarity index 86% rename from java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionTest.java rename to java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitterTest.java index c374f7c139..6e2cccb110 100644 --- a/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionTest.java +++ b/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitterTest.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.completion; +package sleeper.compaction.job.commit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; @@ -42,10 +42,10 @@ import static sleeper.compaction.job.CompactionJobStatusTestData.jobCreated; import static sleeper.core.util.ExponentialBackoffWithJitterTestHelper.noJitter; -public class CompactionJobCompletionTest extends CompactionJobCompletionTestBase { +public class CompactionJobCommitterTest extends CompactionJobCommitterTestBase { @Test - void shouldCompleteCompactionJobsOnDifferentTables() throws Exception { + void shouldCommitCompactionJobsOnDifferentTables() throws Exception { // Given TableProperties table1 = createTable(); TableProperties table2 = createTable(); @@ -59,15 +59,15 @@ void shouldCompleteCompactionJobsOnDifferentTables() throws Exception { RecordsProcessedSummary summary2 = new RecordsProcessedSummary( new RecordsProcessed(450, 400), Instant.parse("2024-05-01T10:58:30Z"), Duration.ofMinutes(1)); - CompactionJobCompletionRequest completion1 = runCompactionJobOnTask("task-1", job1, summary1); - CompactionJobCompletionRequest completion2 = runCompactionJobOnTask("task-2", job2, summary2); + CompactionJobCommitRequest commit1 = runCompactionJobOnTask("task-1", job1, summary1); + CompactionJobCommitRequest commit2 = runCompactionJobOnTask("task-2", job2, summary2); stateStore(table1).fixFileUpdateTime(Instant.parse("2024-05-01T11:00:00Z")); stateStore(table2).fixFileUpdateTime(Instant.parse("2024-05-01T11:00:30Z")); // When - CompactionJobCompletion jobCompletion = jobCompletion(); - jobCompletion.apply(completion1); - jobCompletion.apply(completion2); + CompactionJobCommitter jobCommitter = jobCommitter(); + jobCommitter.apply(commit1); + jobCommitter.apply(commit2); // Then StateStore state1 = stateStore(table1); @@ -102,7 +102,7 @@ void shouldRetryStateStoreUpdateWhenFilesNotAssignedToJob() throws Exception { // Given FileReference file = addInputFile("file.parquet", 123); CompactionJob job = createCompactionJobForOneFileAndRecordStatus(file); - CompactionJobCompletionRequest completion = runCompactionJobOnTask("test-task", job); + CompactionJobCommitRequest commit = runCompactionJobOnTask("test-task", job); actionOnWait(() -> { stateStore().assignJobIds(List.of(AssignJobIdRequest.assignJobOnPartitionToFiles( job.getId(), file.getPartitionId(), List.of(file.getFilename())))); @@ -111,7 +111,7 @@ void shouldRetryStateStoreUpdateWhenFilesNotAssignedToJob() throws Exception { stateStore().fixFileUpdateTime(updateTime); // When - jobCompletion(noJitter()).apply(completion); + jobCommitter(noJitter()).apply(commit); // Then assertThat(stateStore().getFileReferences()).containsExactly( @@ -125,10 +125,10 @@ void shouldFailAfterMaxAttemptsWhenFilesNotAssignedToJob() throws Exception { // Given FileReference file = addInputFile("file.parquet", 123); CompactionJob job = createCompactionJobForOneFileAndRecordStatus(file); - CompactionJobCompletionRequest completion = runCompactionJobOnTask("test-task", job); + CompactionJobCommitRequest commit = runCompactionJobOnTask("test-task", job); // When - assertThatThrownBy(() -> jobCompletion(noJitter()).apply(completion)) + assertThatThrownBy(() -> jobCommitter(noJitter()).apply(commit)) .isInstanceOf(TimedOutWaitingForFileAssignmentsException.class) .hasCauseInstanceOf(FileReferenceNotAssignedToJobException.class); @@ -151,10 +151,10 @@ void shouldFailWithNoRetriesWhenFileDoesNotExistInStateStore() throws Exception // Given FileReference file = inputFileFactory().rootFile("file.parquet", 123); CompactionJob job = createCompactionJobForOneFileAndRecordStatus(file); - CompactionJobCompletionRequest completion = runCompactionJobOnTask("test-task", job); + CompactionJobCommitRequest commit = runCompactionJobOnTask("test-task", job); // When - assertThatThrownBy(() -> jobCompletion(noJitter()).apply(completion)) + assertThatThrownBy(() -> jobCommitter(noJitter()).apply(commit)) .isInstanceOf(FileNotFoundException.class); // Then diff --git a/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionTestBase.java b/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitterTestBase.java similarity index 89% rename from java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionTestBase.java rename to java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitterTestBase.java index 73f48bfcb7..15573aa55d 100644 --- a/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/completion/CompactionJobCompletionTestBase.java +++ b/java/compaction/compaction-core/src/test/java/sleeper/compaction/job/commit/CompactionJobCommitterTestBase.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package sleeper.compaction.job.completion; +package sleeper.compaction.job.commit; import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionJobFactory; @@ -49,7 +49,7 @@ import static sleeper.core.util.ExponentialBackoffWithJitterTestHelper.fixJitterSeed; import static sleeper.core.util.ExponentialBackoffWithJitterTestHelper.recordWaits; -public class CompactionJobCompletionTestBase { +public class CompactionJobCommitterTestBase { private static final Instant INPUT_UPDATE_TIME = Instant.parse("2024-05-01T10:00:00Z"); private static final RecordsProcessedSummary DEFAULT_SUMMARY = new RecordsProcessedSummary( @@ -106,22 +106,22 @@ protected CompactionJob createCompactionJobForOneFileAndRecordStatus(TableProper return job; } - protected CompactionJobCompletionRequest runCompactionJobOnTask(String taskId, CompactionJob job) throws Exception { + protected CompactionJobCommitRequest runCompactionJobOnTask(String taskId, CompactionJob job) throws Exception { return runCompactionJobOnTask(taskId, job, DEFAULT_SUMMARY); } - protected CompactionJobCompletionRequest runCompactionJobOnTask(String taskId, CompactionJob job, RecordsProcessedSummary summary) throws Exception { + protected CompactionJobCommitRequest runCompactionJobOnTask(String taskId, CompactionJob job, RecordsProcessedSummary summary) throws Exception { statusStore.jobStarted(job, summary.getStartTime(), taskId); - return new CompactionJobCompletionRequest(job, taskId, summary); + return new CompactionJobCommitRequest(job, taskId, summary); } - protected CompactionJobCompletion jobCompletion() { - return jobCompletion(fixJitterSeed()); + protected CompactionJobCommitter jobCommitter() { + return jobCommitter(fixJitterSeed()); } - protected CompactionJobCompletion jobCompletion(DoubleSupplier randomJitter) { - return new CompactionJobCompletion(statusStore, stateStoreByTableId::get, - CompactionJobCompletion.JOB_ASSIGNMENT_WAIT_ATTEMPTS, backoff(randomJitter)); + protected CompactionJobCommitter jobCommitter(DoubleSupplier randomJitter) { + return new CompactionJobCommitter(statusStore, stateStoreByTableId::get, + CompactionJobCommitter.JOB_ASSIGNMENT_WAIT_ATTEMPTS, backoff(randomJitter)); } protected FileReferenceFactory fileFactory(TableProperties table, Instant updateTime) { @@ -155,7 +155,7 @@ protected void actionOnWait(WaitAction action) throws Exception { private ExponentialBackoffWithJitter backoff(DoubleSupplier randomJitter) { return new ExponentialBackoffWithJitter( - CompactionJobCompletion.JOB_ASSIGNMENT_WAIT_RANGE, + CompactionJobCommitter.JOB_ASSIGNMENT_WAIT_RANGE, randomJitter, waiter); } diff --git a/java/compaction/compaction-job-completion-lambda/pom.xml b/java/compaction/compaction-job-committer-lambda/pom.xml similarity index 98% rename from java/compaction/compaction-job-completion-lambda/pom.xml rename to java/compaction/compaction-job-committer-lambda/pom.xml index 72dd51967e..5199c3218d 100644 --- a/java/compaction/compaction-job-completion-lambda/pom.xml +++ b/java/compaction/compaction-job-committer-lambda/pom.xml @@ -23,7 +23,7 @@ 4.0.0 - compaction-job-completion-lambda + compaction-job-committer-lambda diff --git a/java/compaction/compaction-job-completion-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCompletionLambda.java b/java/compaction/compaction-job-committer-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCommitterLambda.java similarity index 75% rename from java/compaction/compaction-job-completion-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCompletionLambda.java rename to java/compaction/compaction-job-committer-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCommitterLambda.java index 435623bc2f..19c9a57192 100644 --- a/java/compaction/compaction-job-completion-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCompletionLambda.java +++ b/java/compaction/compaction-job-committer-lambda/src/main/java/sleeper/compaction/completion/lambda/CompactionJobCommitterLambda.java @@ -30,10 +30,10 @@ import org.slf4j.LoggerFactory; import sleeper.compaction.job.CompactionJobStatusStore; -import sleeper.compaction.job.completion.CompactionJobCompletion; -import sleeper.compaction.job.completion.CompactionJobCompletion.GetStateStore; -import sleeper.compaction.job.completion.CompactionJobCompletionRequest; -import sleeper.compaction.job.completion.CompactionJobCompletionRequestSerDe; +import sleeper.compaction.job.commit.CompactionJobCommitRequest; +import sleeper.compaction.job.commit.CompactionJobCommitRequestSerDe; +import sleeper.compaction.job.commit.CompactionJobCommitter; +import sleeper.compaction.job.commit.CompactionJobCommitter.GetStateStore; import sleeper.compaction.status.store.job.CompactionJobStatusStoreFactory; import sleeper.configuration.properties.instance.InstanceProperties; import sleeper.configuration.properties.table.TablePropertiesProvider; @@ -48,18 +48,18 @@ import static sleeper.configuration.properties.instance.CdkDefinedInstanceProperty.CONFIG_BUCKET; -public class CompactionJobCompletionLambda implements RequestHandler { - public static final Logger LOGGER = LoggerFactory.getLogger(CompactionJobCompletionLambda.class); +public class CompactionJobCommitterLambda implements RequestHandler { + public static final Logger LOGGER = LoggerFactory.getLogger(CompactionJobCommitterLambda.class); - private final CompactionJobCompletion compactionJobCompletion; - private final CompactionJobCompletionRequestSerDe serDe = new CompactionJobCompletionRequestSerDe(); + private final CompactionJobCommitter compactionJobCommitter; + private final CompactionJobCommitRequestSerDe serDe = new CompactionJobCommitRequestSerDe(); - public CompactionJobCompletionLambda() { + public CompactionJobCommitterLambda() { this(connectToAws()); } - public CompactionJobCompletionLambda(CompactionJobCompletion compactionJobCompletion) { - this.compactionJobCompletion = compactionJobCompletion; + public CompactionJobCommitterLambda(CompactionJobCommitter compactionJobCommitter) { + this.compactionJobCommitter = compactionJobCommitter; } @Override @@ -70,11 +70,11 @@ public SQSBatchResponse handleRequest(SQSEvent event, Context context) { for (SQSMessage message : event.getRecords()) { try { LOGGER.info("Found message: {}", message.getBody()); - CompactionJobCompletionRequest request = serDe.fromJson(message.getBody()); - compactionJobCompletion.apply(request); - LOGGER.info("Completed"); + CompactionJobCommitRequest request = serDe.fromJson(message.getBody()); + compactionJobCommitter.apply(request); + LOGGER.info("Successfully committed compaction job {}", request.getJob()); } catch (RuntimeException | StateStoreException | InterruptedException e) { - LOGGER.error("Failed completing compaction job", e); + LOGGER.error("Failed committing compaction job", e); batchItemFailures.add(new BatchItemFailure(message.getMessageId())); } } @@ -84,7 +84,7 @@ public SQSBatchResponse handleRequest(SQSEvent event, Context context) { return new SQSBatchResponse(batchItemFailures); } - private static CompactionJobCompletion connectToAws() { + private static CompactionJobCommitter connectToAws() { AmazonS3 s3Client = AmazonS3ClientBuilder.defaultClient(); AmazonDynamoDB dynamoDBClient = AmazonDynamoDBClientBuilder.defaultClient(); String s3Bucket = System.getenv(CONFIG_BUCKET.toEnvironmentVariable()); @@ -96,11 +96,11 @@ private static CompactionJobCompletion connectToAws() { TablePropertiesProvider tablePropertiesProvider = new TablePropertiesProvider(instanceProperties, s3Client, dynamoDBClient); StateStoreProvider stateStoreProvider = new StateStoreProvider(instanceProperties, s3Client, dynamoDBClient, hadoopConf); CompactionJobStatusStore statusStore = CompactionJobStatusStoreFactory.getStatusStore(dynamoDBClient, instanceProperties); - return new CompactionJobCompletion( - statusStore, stateStoreProviderForCompletion(tablePropertiesProvider, stateStoreProvider)); + return new CompactionJobCommitter( + statusStore, stateStoreProviderForCommitter(tablePropertiesProvider, stateStoreProvider)); } - private static GetStateStore stateStoreProviderForCompletion( + private static GetStateStore stateStoreProviderForCommitter( TablePropertiesProvider tablePropertiesProvider, StateStoreProvider stateStoreProvider) { return tableId -> stateStoreProvider.getStateStore(tablePropertiesProvider.getById(tableId)); } diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionJobCommitHandler.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionJobCommitHandler.java new file mode 100644 index 0000000000..9d9b1ff250 --- /dev/null +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionJobCommitHandler.java @@ -0,0 +1,55 @@ +/* + * 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.compaction.job.execution; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import sleeper.compaction.job.commit.CompactionJobCommitRequest; +import sleeper.compaction.job.commit.CompactionJobCommitter; +import sleeper.configuration.properties.table.TablePropertiesProvider; +import sleeper.core.statestore.StateStoreException; + +import static sleeper.configuration.properties.table.TableProperty.COMPACTION_JOB_COMMIT_ASYNC; + +public class CompactionJobCommitHandler { + public static final Logger LOGGER = LoggerFactory.getLogger(CompactionJobCommitHandler.class); + + private TablePropertiesProvider tablePropertiesProvider; + private CompactionJobCommitter jobCommitter; + private CommitQueueSender jobCommitQueueSender; + + public CompactionJobCommitHandler(TablePropertiesProvider tablePropertiesProvider, + CompactionJobCommitter jobCommitter, CommitQueueSender jobCommitQueueSender) { + this.tablePropertiesProvider = tablePropertiesProvider; + this.jobCommitter = jobCommitter; + this.jobCommitQueueSender = jobCommitQueueSender; + } + + public void commit(CompactionJobCommitRequest commitRequest) throws StateStoreException, InterruptedException { + if (tablePropertiesProvider.getById(commitRequest.getJob().getTableId()).getBoolean(COMPACTION_JOB_COMMIT_ASYNC)) { + LOGGER.info("Sending compaction job {} to queue to be committed asynchronously", commitRequest.getJob().getId()); + jobCommitQueueSender.send(commitRequest); + } else { + LOGGER.info("Committing compaction job {} inside compaction task", commitRequest.getJob().getId()); + jobCommitter.apply(commitRequest); + } + } + + interface CommitQueueSender { + void send(CompactionJobCommitRequest commitRequest); + } +} diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionTask.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionTask.java index b7dab4f6c7..8d4a0a7056 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionTask.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/CompactionTask.java @@ -22,6 +22,7 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionJobStatusStore; import sleeper.compaction.job.CompactionRunner; +import sleeper.compaction.job.commit.CompactionJobCommitRequest; import sleeper.compaction.task.CompactionTaskFinishedStatus; import sleeper.compaction.task.CompactionTaskStatus; import sleeper.compaction.task.CompactionTaskStatusStore; @@ -62,16 +63,18 @@ public class CompactionTask { private final PropertiesReloader propertiesReloader; private int numConsecutiveFailures = 0; private int totalNumberOfMessagesProcessed = 0; + private CompactionJobCommitHandler jobCommitHandler; public CompactionTask(InstanceProperties instanceProperties, PropertiesReloader propertiesReloader, - MessageReceiver messageReceiver, CompactionAlgorithmSelector selector, - CompactionJobStatusStore jobStore, CompactionTaskStatusStore taskStore, String taskId) { - this(instanceProperties, propertiesReloader, messageReceiver, selector, jobStore, taskStore, taskId, Instant::now, threadSleep()); + MessageReceiver messageReceiver, CompactionRunner compactor, CompactionJobCommitHandler jobCommitHandler, + CompactionJobStatusStore jobStore, CompactionTaskStatusStore taskStore, CompactionAlgorithmSelector selector, String taskId) { + this(instanceProperties, propertiesReloader, messageReceiver, compactor, jobCommitHandler, jobStore, taskStore, selector, taskId, Instant::now, threadSleep()); } public CompactionTask(InstanceProperties instanceProperties, PropertiesReloader propertiesReloader, - MessageReceiver messageReceiver, CompactionAlgorithmSelector selector, CompactionJobStatusStore jobStore, - CompactionTaskStatusStore taskStore, String taskId, Supplier timeSupplier, Consumer sleepForTime) { + MessageReceiver messageReceiver, CompactionRunner compactor, CompactionJobCommitHandler jobCommitHandler, + CompactionJobStatusStore jobStore, CompactionTaskStatusStore taskStore, CompactionAlgorithmSelector selector, String taskId, Supplier timeSupplier, + Consumer sleepForTime) { maxIdleTime = Duration.ofSeconds(instanceProperties.getInt(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS)); maxConsecutiveFailures = instanceProperties.getInt(COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES); delayBeforeRetry = Duration.ofSeconds(instanceProperties.getInt(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS)); @@ -83,6 +86,7 @@ public CompactionTask(InstanceProperties instanceProperties, PropertiesReloader this.jobStatusStore = jobStore; this.taskStatusStore = taskStore; this.taskId = taskId; + this.jobCommitHandler = jobCommitHandler; } public void run() throws IOException { @@ -152,7 +156,7 @@ private RecordsProcessedSummary compact(CompactionJob job) throws Exception { RecordsProcessed recordsProcessed = compactor.compact(job); Instant jobFinishTime = timeSupplier.get(); RecordsProcessedSummary summary = new RecordsProcessedSummary(recordsProcessed, jobStartTime, jobFinishTime); - jobStatusStore.jobFinished(job, summary, taskId); + jobCommitHandler.commit(new CompactionJobCommitRequest(job, taskId, summary)); logMetrics(job, summary); return summary; } diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/ECSCompactionTaskRunner.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/ECSCompactionTaskRunner.java index 7278f68f7e..e1f4f93147 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/ECSCompactionTaskRunner.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/ECSCompactionTaskRunner.java @@ -28,6 +28,7 @@ import org.slf4j.LoggerFactory; import sleeper.compaction.job.CompactionJobStatusStore; +import sleeper.compaction.job.commit.CompactionJobCommitter; import sleeper.compaction.status.store.job.CompactionJobStatusStoreFactory; import sleeper.compaction.status.store.task.CompactionTaskStatusStoreFactory; import sleeper.compaction.task.CompactionTaskStatusStore; @@ -36,6 +37,7 @@ import sleeper.configuration.properties.PropertiesReloader; import sleeper.configuration.properties.instance.InstanceProperties; import sleeper.configuration.properties.table.TablePropertiesProvider; +import sleeper.core.statestore.StateStoreException; import sleeper.core.util.LoggedDuration; import sleeper.io.parquet.utils.HadoopConfigurationProvider; import sleeper.job.common.EC2ContainerMetadata; @@ -93,9 +95,19 @@ public static void main(String[] args) throws IOException, ObjectFactoryExceptio DefaultSelector compactionSelector = new DefaultSelector(instanceProperties, tablePropertiesProvider, stateStoreProvider, objectFactory); + CompactionJobCommitter committer = new CompactionJobCommitter(jobStatusStore, tableId -> stateStoreProvider.getStateStore(tablePropertiesProvider.getById(tableId))); + CompactionJobCommitHandler commitHandler = new CompactionJobCommitHandler(tablePropertiesProvider, committer, + (request) -> { + // TODO send to SQS queue once infrastructure is deployed by CDK + try { + committer.apply(request); + } catch (StateStoreException | InterruptedException e) { + throw new RuntimeException(e); + } + }); CompactionTask task = new CompactionTask(instanceProperties, propertiesReloader, - new SqsCompactionQueueHandler(sqsClient, instanceProperties), - compactionSelector, jobStatusStore, taskStatusStore, taskId); + new SqsCompactionQueueHandler(sqsClient, instanceProperties), commitHandler, jobStatusStore, + taskStatusStore, compactionSelector, taskId); task.run(); } finally { sqsClient.shutdown(); diff --git a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java index a8e3285195..51fa12f831 100644 --- a/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java +++ b/java/compaction/compaction-job-execution/src/main/java/sleeper/compaction/job/execution/StandardCompactor.java @@ -26,7 +26,6 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionRunner; -import sleeper.compaction.job.StateStoreUpdate; import sleeper.configuration.jars.ObjectFactory; import sleeper.configuration.jars.ObjectFactoryException; import sleeper.configuration.properties.instance.InstanceProperties; @@ -52,7 +51,6 @@ import sleeper.statestore.StateStoreProvider; import java.io.IOException; -import java.time.LocalDateTime; import java.util.ArrayList; import java.util.List; import java.util.NoSuchElementException; @@ -137,17 +135,9 @@ public RecordsProcessed compact(CompactionJob compactionJob) throws IOException, } LOGGER.info("Compaction job {}: Read {} records and wrote {} records", compactionJob.getId(), totalNumberOfRecordsRead, recordsWritten); - - StateStoreUpdate.updateStateStoreSuccess(compactionJob, recordsWritten, stateStore); - LOGGER.info("Compaction job {}: compaction committed to state store at {}", compactionJob.getId(), LocalDateTime.now()); - return new RecordsProcessed(totalNumberOfRecordsRead, recordsWritten); } - private Configuration getConfiguration() { - return HadoopConfigurationProvider.getConfigurationForECS(instanceProperties); - } - private List> createInputIterators(CompactionJob compactionJob, Partition partition, Schema schema, Configuration conf) throws IOException { List> inputIterators = new ArrayList<>(); @@ -187,4 +177,8 @@ public static CloseableIterator getMergingIterator( } return mergingIterator; } + + private Configuration getConfiguration() { + return HadoopConfigurationProvider.getConfigurationForECS(instanceProperties); + } } diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesEmptyOutputIT.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesEmptyOutputIT.java index e0a7a3b566..1c63b76022 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesEmptyOutputIT.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesEmptyOutputIT.java @@ -26,9 +26,7 @@ import sleeper.core.schema.Schema; import sleeper.core.schema.type.LongType; import sleeper.core.statestore.FileReference; -import sleeper.core.statestore.FileReferenceFactory; -import java.time.Instant; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -64,16 +62,6 @@ void shouldMergeFilesCorrectlyWhenSomeAreEmpty() throws Exception { assertThat(summary.getRecordsRead()).isEqualTo(data.size()); assertThat(summary.getRecordsWritten()).isEqualTo(data.size()); assertThat(readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(data); - - // - Check state store has correct ready for GC files - assertThat(stateStore.getReadyForGCFilenamesBefore(Instant.ofEpochMilli(Long.MAX_VALUE))) - .containsExactlyInAnyOrder(file1.getFilename(), file2.getFilename()); - - // - Check state store has correct file references - assertThat(stateStore.getFileReferences()) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields("lastStateStoreUpdateTime") - .containsExactly(FileReferenceFactory.from(stateStore) - .rootFile(compactionJob.getOutputFile(), 100L)); } @Test @@ -99,15 +87,5 @@ void shouldMergeFilesCorrectlyWhenAllAreEmpty() throws Exception { assertThat(summary.getRecordsRead()).isZero(); assertThat(summary.getRecordsWritten()).isZero(); assertThat(readDataFile(schema, compactionJob.getOutputFile())).isEmpty(); - - // - Check state store has correct ready for GC files - assertThat(stateStore.getReadyForGCFilenamesBefore(Instant.ofEpochMilli(Long.MAX_VALUE))) - .containsExactly(file1.getFilename(), file2.getFilename()); - - // - Check state store has correct active files - assertThat(stateStore.getFileReferences()) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields("lastStateStoreUpdateTime") - .containsExactly(FileReferenceFactory.from(stateStore) - .rootFile(compactionJob.getOutputFile(), 0L)); } } diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIT.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIT.java index e8592265bb..b0664a27bd 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIT.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIT.java @@ -15,7 +15,6 @@ */ package sleeper.compaction.job.execution; -import org.assertj.core.api.Assertions; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; @@ -32,9 +31,7 @@ import sleeper.core.schema.type.LongType; import sleeper.core.schema.type.StringType; import sleeper.core.statestore.FileReference; -import sleeper.core.statestore.FileReferenceFactory; -import java.time.Instant; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -68,17 +65,7 @@ void shouldMergeFilesCorrectlyAndUpdateStateStoreWithLongKey() throws Exception List expectedResults = CompactSortedFilesTestData.combineSortedBySingleKey(data1, data2); assertThat(summary.getRecordsRead()).isEqualTo(expectedResults.size()); assertThat(summary.getRecordsWritten()).isEqualTo(expectedResults.size()); - Assertions.assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); - - // - Check DynamoDBStateStore has correct ready for GC files - assertThat(stateStore.getReadyForGCFilenamesBefore(Instant.ofEpochMilli(Long.MAX_VALUE))) - .containsExactlyInAnyOrder(file1.getFilename(), file2.getFilename()); - - // - Check DynamoDBStateStore has correct file references - assertThat(stateStore.getFileReferences()) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields("lastStateStoreUpdateTime") - .containsExactly(FileReferenceFactory.from(stateStore) - .rootFile(compactionJob.getOutputFile(), 200L)); + assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); } @Nested @@ -126,17 +113,7 @@ void shouldMergeFilesCorrectlyAndUpdateStateStoreWithStringKey() throws Exceptio List expectedResults = CompactSortedFilesTestData.combineSortedBySingleKey(data1, data2); assertThat(summary.getRecordsRead()).isEqualTo(expectedResults.size()); assertThat(summary.getRecordsWritten()).isEqualTo(expectedResults.size()); - Assertions.assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); - - // - Check DynamoDBStateStore has correct ready for GC files - assertThat(stateStore.getReadyForGCFilenamesBefore(Instant.ofEpochMilli(Long.MAX_VALUE))) - .containsExactlyInAnyOrder(file1.getFilename(), file2.getFilename()); - - // - Check DynamoDBStateStore has correct file references - assertThat(stateStore.getFileReferences()) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields("lastStateStoreUpdateTime") - .containsExactly(FileReferenceFactory.from(stateStore) - .rootFile(compactionJob.getOutputFile(), 200L)); + assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); } } @@ -191,17 +168,7 @@ void shouldMergeFilesCorrectlyAndUpdateStateStoreWithByteArrayKey() throws Excep List expectedResults = CompactSortedFilesTestData.combineSortedBySingleByteArrayKey(data1, data2); assertThat(summary.getRecordsRead()).isEqualTo(expectedResults.size()); assertThat(summary.getRecordsWritten()).isEqualTo(expectedResults.size()); - Assertions.assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); - - // - Check DynamoDBStateStore has correct ready for GC files - assertThat(stateStore.getReadyForGCFilenamesBefore(Instant.ofEpochMilli(Long.MAX_VALUE))) - .containsExactlyInAnyOrder(file1.getFilename(), file2.getFilename()); - - // - Check DynamoDBStateStore has correct file references - assertThat(stateStore.getFileReferences()) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields("lastStateStoreUpdateTime") - .containsExactly(FileReferenceFactory.from(stateStore) - .rootFile(compactionJob.getOutputFile(), 200L)); + assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); } } } diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIteratorIT.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIteratorIT.java index 4b9f1ecc2d..69934dc2c4 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIteratorIT.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesIteratorIT.java @@ -15,7 +15,6 @@ */ package sleeper.compaction.job.execution; -import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; import sleeper.compaction.job.CompactionJob; @@ -29,9 +28,7 @@ import sleeper.core.record.process.RecordsProcessed; import sleeper.core.schema.Schema; import sleeper.core.statestore.FileReference; -import sleeper.core.statestore.FileReferenceFactory; -import java.time.Instant; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -75,16 +72,6 @@ void shouldApplyIteratorDuringCompaction() throws Exception { // - Read output files and check that they contain the right results assertThat(summary.getRecordsRead()).isEqualTo(200L); assertThat(summary.getRecordsWritten()).isEqualTo(100L); - Assertions.assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(data1); - - // - Check DynamoDBStateStore has correct ready for GC files - assertThat(stateStore.getReadyForGCFilenamesBefore(Instant.ofEpochMilli(Long.MAX_VALUE))) - .containsExactlyInAnyOrder(file1.getFilename(), file2.getFilename()); - - // - Check DynamoDBStateStore has correct file references - assertThat(stateStore.getFileReferences()) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields("lastStateStoreUpdateTime") - .containsExactly(FileReferenceFactory.from(stateStore) - .rootFile(compactionJob.getOutputFile(), 100L)); + assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(data1); } } diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesLocalStackIT.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesLocalStackIT.java index 180b3ddb8e..84d34a5d68 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesLocalStackIT.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesLocalStackIT.java @@ -19,7 +19,6 @@ import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder; import com.amazonaws.services.s3.AmazonS3; import com.amazonaws.services.s3.AmazonS3ClientBuilder; -import org.assertj.core.api.Assertions; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; @@ -46,13 +45,11 @@ import sleeper.core.schema.Schema; import sleeper.core.schema.type.LongType; import sleeper.core.statestore.FileReference; -import sleeper.core.statestore.FileReferenceFactory; import sleeper.core.statestore.StateStore; import sleeper.statestore.FixedStateStoreProvider; import sleeper.statestore.StateStoreFactory; import sleeper.statestore.s3.S3StateStoreCreator; -import java.time.Instant; import java.util.List; import static org.assertj.core.api.Assertions.assertThat; @@ -149,15 +146,6 @@ public void shouldUpdateStateStoreAfterRunningCompactionJob() throws Exception { List expectedResults = CompactSortedFilesTestData.combineSortedBySingleKey(data1, data2); assertThat(summary.getRecordsRead()).isEqualTo(expectedResults.size()); assertThat(summary.getRecordsWritten()).isEqualTo(expectedResults.size()); - Assertions.assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); - - // - Check StateStore has correct ready for GC files - assertThat(stateStore.getReadyForGCFilenamesBefore(Instant.ofEpochMilli(Long.MAX_VALUE))) - .containsExactlyInAnyOrder(file1.getFilename(), file2.getFilename()); - - // - Check StateStore has correct file references - assertThat(stateStore.getFileReferences()) - .usingRecursiveFieldByFieldElementComparatorIgnoringFields("lastStateStoreUpdateTime") - .containsExactly(FileReferenceFactory.from(tree).rootFile(compactionJob.getOutputFile(), 200L)); + assertThat(CompactSortedFilesTestData.readDataFile(schema, compactionJob.getOutputFile())).isEqualTo(expectedResults); } } diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java index c3bb3fa596..5dbb824f87 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactSortedFilesRetryStateStoreTest.java @@ -21,7 +21,7 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionJobFactory; import sleeper.compaction.job.StateStoreUpdate; -import sleeper.compaction.job.completion.TimedOutWaitingForFileAssignmentsException; +import sleeper.compaction.job.commit.TimedOutWaitingForFileAssignmentsException; import sleeper.configuration.properties.instance.InstanceProperties; import sleeper.configuration.properties.table.TableProperties; import sleeper.core.partition.PartitionTree; @@ -44,6 +44,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static sleeper.compaction.job.commit.CompactionJobCommitterUtils.updateStateStoreSuccess; import static sleeper.configuration.properties.InstancePropertiesTestHelper.createTestInstanceProperties; import static sleeper.configuration.properties.table.TablePropertiesTestHelper.createTestTableProperties; import static sleeper.core.schema.SchemaTestHelper.schemaWithKey; @@ -82,7 +83,7 @@ void shouldRetryStateStoreUpdateWhenFilesNotAssignedToJob() throws Exception { }); // When - updateStateStoreSuccess(job, 123, noJitter()); + updateStateStore(job, 123, noJitter()); // Then assertThat(stateStore.getFileReferences()).containsExactly( @@ -98,7 +99,7 @@ void shouldFailAfterMaxAttemptsWhenFilesNotAssignedToJob() throws Exception { CompactionJob job = createCompactionJobForOneFile(file); // When - assertThatThrownBy(() -> updateStateStoreSuccess(job, 123, noJitter())) + assertThatThrownBy(() -> updateStateStore(job, 123, noJitter())) .isInstanceOf(TimedOutWaitingForFileAssignmentsException.class) .hasCauseInstanceOf(FileReferenceNotAssignedToJobException.class); @@ -123,7 +124,7 @@ void shouldFailWithNoRetriesWhenFileDoesNotExist() throws Exception { CompactionJob job = createCompactionJobForOneFile(file); // When - assertThatThrownBy(() -> updateStateStoreSuccess(job, 123, noJitter())) + assertThatThrownBy(() -> updateStateStore(job, 123, noJitter())) .isInstanceOf(FileNotFoundException.class); // Then @@ -131,8 +132,8 @@ void shouldFailWithNoRetriesWhenFileDoesNotExist() throws Exception { assertThat(foundWaits).isEmpty(); } - private void updateStateStoreSuccess(CompactionJob job, long recordsWritten, DoubleSupplier randomJitter) throws Exception { - StateStoreUpdate.updateStateStoreSuccess(job, 123, stateStore, + private void updateStateStore(CompactionJob job, long recordsWritten, DoubleSupplier randomJitter) throws Exception { + updateStateStoreSuccess(job, 123, stateStore, StateStoreUpdate.JOB_ASSIGNMENT_WAIT_ATTEMPTS, backoff(randomJitter)); } diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskCommitTest.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskCommitTest.java new file mode 100644 index 0000000000..79808cce84 --- /dev/null +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskCommitTest.java @@ -0,0 +1,313 @@ +/* + * 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.compaction.job.execution; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +import sleeper.compaction.job.CompactionJob; +import sleeper.compaction.job.commit.CompactionJobCommitRequest; +import sleeper.compaction.task.CompactionTaskFinishedStatus; +import sleeper.compaction.task.CompactionTaskStatus; +import sleeper.configuration.properties.table.FixedTablePropertiesProvider; +import sleeper.configuration.properties.table.TableProperties; +import sleeper.core.record.process.RecordsProcessed; +import sleeper.core.record.process.RecordsProcessedSummary; +import sleeper.core.statestore.StateStore; +import sleeper.statestore.FixedStateStoreProvider; + +import java.time.Instant; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Queue; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static sleeper.compaction.job.CompactionJobStatusTestData.finishedCompactionRun; +import static sleeper.compaction.job.CompactionJobStatusTestData.jobCreated; +import static sleeper.compaction.job.CompactionJobStatusTestData.startedCompactionRun; +import static sleeper.configuration.properties.instance.DefaultProperty.DEFAULT_COMPACTION_JOB_COMMIT_ASYNC; +import static sleeper.configuration.properties.table.TableProperty.COMPACTION_JOB_COMMIT_ASYNC; +import static sleeper.configuration.properties.table.TableProperty.TABLE_ID; +import static sleeper.core.statestore.inmemory.StateStoreTestHelper.inMemoryStateStoreWithSinglePartition; + +public class CompactionTaskCommitTest extends CompactionTaskTestBase { + private final TableProperties table1 = createTable("test-table-1-id", "test-table-1"); + private final TableProperties table2 = createTable("test-table-2-id", "test-table-2"); + private final StateStore store1 = inMemoryStateStoreWithSinglePartition(schema); + private final StateStore store2 = inMemoryStateStoreWithSinglePartition(schema); + private final FixedTablePropertiesProvider tablePropertiesProvider = new FixedTablePropertiesProvider(List.of(table1, table2)); + private final FixedStateStoreProvider stateStoreProvider = new FixedStateStoreProvider(Map.of("test-table-1", store1, "test-table-2", store2)); + + @Nested + @DisplayName("Send commits to queue") + class SendCommitsToQueue { + @BeforeEach + public void setup() { + instanceProperties.set(DEFAULT_COMPACTION_JOB_COMMIT_ASYNC, "true"); + } + + @Test + void shouldSendJobCommitRequestToQueue() throws Exception { + // Given + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // Job started + Instant.parse("2024-02-22T13:50:02Z"), // Job completed + Instant.parse("2024-02-22T13:50:05Z"))); // Finish + CompactionJob job1 = createJobOnQueue("job1"); + RecordsProcessed job1Summary = new RecordsProcessed(10L, 5L); + + // When + runTask(processJobs(jobSucceeds(job1Summary)), times::poll); + + // Then + assertThat(successfulJobs).containsExactly(job1); + assertThat(failedJobs).isEmpty(); + assertThat(jobsOnQueue).isEmpty(); + assertThat(commitRequestsOnQueue).containsExactly( + commitRequestFor(job1, + new RecordsProcessedSummary(job1Summary, + Instant.parse("2024-02-22T13:50:01Z"), + Instant.parse("2024-02-22T13:50:02Z")))); + assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).containsExactly( + jobCreated(job1, DEFAULT_CREATED_TIME, + startedCompactionRun(DEFAULT_TASK_ID, + Instant.parse("2024-02-22T13:50:01Z")))); + } + + @Test + void shouldSendJobCommitRequestsForDifferentTablesToQueue() throws Exception { + // Given + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // Job 1 started + Instant.parse("2024-02-22T13:50:02Z"), // Job 1 completed + Instant.parse("2024-02-22T13:50:03Z"), // Job 2 started + Instant.parse("2024-02-22T13:50:04Z"), // Job 2 completed + Instant.parse("2024-02-22T13:50:07Z"))); // Finish + CompactionJob job1 = createJobOnQueue("job1", table1, store1); + RecordsProcessed job1Summary = new RecordsProcessed(10L, 10L); + CompactionJob job2 = createJobOnQueue("job2", table2, store2); + RecordsProcessed job2Summary = new RecordsProcessed(20L, 20L); + + // When + runTask(processJobs( + jobSucceeds(job1Summary), + jobSucceeds(job2Summary)), + times::poll, tablePropertiesProvider, stateStoreProvider); + + // Then + assertThat(successfulJobs).containsExactly(job1, job2); + assertThat(failedJobs).isEmpty(); + assertThat(jobsOnQueue).isEmpty(); + assertThat(commitRequestsOnQueue).containsExactly( + commitRequestFor(job1, + new RecordsProcessedSummary(job1Summary, + Instant.parse("2024-02-22T13:50:01Z"), + Instant.parse("2024-02-22T13:50:02Z"))), + commitRequestFor(job2, + new RecordsProcessedSummary(job2Summary, + Instant.parse("2024-02-22T13:50:03Z"), + Instant.parse("2024-02-22T13:50:04Z")))); + assertThat(jobStore.getAllJobs(table1.get(TABLE_ID))).containsExactly( + jobCreated(job1, DEFAULT_CREATED_TIME, + startedCompactionRun(DEFAULT_TASK_ID, + Instant.parse("2024-02-22T13:50:01Z")))); + assertThat(jobStore.getAllJobs(table2.get(TABLE_ID))).containsExactly( + jobCreated(job2, DEFAULT_CREATED_TIME, + startedCompactionRun(DEFAULT_TASK_ID, + Instant.parse("2024-02-22T13:50:03Z")))); + } + + @Test + void shouldOnlySendJobCommitRequestsForTablesConfiguredForAsyncCommit() throws Exception { + // Given + table2.set(COMPACTION_JOB_COMMIT_ASYNC, "false"); + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // Job 1 started + Instant.parse("2024-02-22T13:50:02Z"), // Job 1 completed + Instant.parse("2024-02-22T13:50:03Z"), // Job 2 started + Instant.parse("2024-02-22T13:50:04Z"), // Job 2 completed + Instant.parse("2024-02-22T13:50:07Z"))); // Finish + CompactionJob job1 = createJobOnQueue("job1", table1, store1); + RecordsProcessed job1Summary = new RecordsProcessed(10L, 10L); + CompactionJob job2 = createJobOnQueue("job2", table2, store2); + RecordsProcessed job2Summary = new RecordsProcessed(20L, 20L); + + // When + runTask(processJobs( + jobSucceeds(job1Summary), + jobSucceeds(job2Summary)), + times::poll, tablePropertiesProvider, stateStoreProvider); + + // Then + assertThat(successfulJobs).containsExactly(job1, job2); + assertThat(failedJobs).isEmpty(); + assertThat(jobsOnQueue).isEmpty(); + assertThat(commitRequestsOnQueue).containsExactly( + commitRequestFor(job1, + new RecordsProcessedSummary(job1Summary, + Instant.parse("2024-02-22T13:50:01Z"), + Instant.parse("2024-02-22T13:50:02Z")))); + assertThat(jobStore.getAllJobs(table1.get(TABLE_ID))).containsExactly( + jobCreated(job1, DEFAULT_CREATED_TIME, + startedCompactionRun(DEFAULT_TASK_ID, + Instant.parse("2024-02-22T13:50:01Z")))); + assertThat(jobStore.getAllJobs(table2.get(TABLE_ID))).containsExactly( + jobCreated(job2, DEFAULT_CREATED_TIME, + finishedCompactionRun(DEFAULT_TASK_ID, new RecordsProcessedSummary(job2Summary, + Instant.parse("2024-02-22T13:50:03Z"), + Instant.parse("2024-02-22T13:50:04Z"))))); + } + + private CompactionJobCommitRequest commitRequestFor(CompactionJob job, RecordsProcessedSummary summary) { + return new CompactionJobCommitRequest(job, DEFAULT_TASK_ID, summary); + } + } + + @Nested + @DisplayName("Update status stores") + class UpdateStatusStores { + @Test + void shouldSaveTaskAndJobWhenOneJobSucceeds() throws Exception { + // Given + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // Job started + Instant.parse("2024-02-22T13:50:02Z"), // Job completed + Instant.parse("2024-02-22T13:50:05Z"))); // Finish + CompactionJob job = createJobOnQueue("job1"); + + // When + RecordsProcessed recordsProcessed = new RecordsProcessed(10L, 10L); + runTask("test-task-1", processJobs( + jobSucceeds(recordsProcessed)), + times::poll); + + // Then + RecordsProcessedSummary jobSummary = new RecordsProcessedSummary(recordsProcessed, + Instant.parse("2024-02-22T13:50:01Z"), + Instant.parse("2024-02-22T13:50:02Z")); + assertThat(taskStore.getAllTasks()).containsExactly( + finishedCompactionTask("test-task-1", + Instant.parse("2024-02-22T13:50:00Z"), + Instant.parse("2024-02-22T13:50:05Z"), + jobSummary)); + assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).containsExactly( + jobCreated(job, DEFAULT_CREATED_TIME, + finishedCompactionRun("test-task-1", jobSummary))); + } + + @Test + void shouldSaveTaskAndJobsWhenMultipleJobsSucceed() throws Exception { + // Given + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // Job 1 started + Instant.parse("2024-02-22T13:50:02Z"), // Job 1 completed + Instant.parse("2024-02-22T13:50:03Z"), // Job 2 started + Instant.parse("2024-02-22T13:50:04Z"), // Job 2 completed + Instant.parse("2024-02-22T13:50:05Z"))); // Finish + CompactionJob job1 = createJobOnQueue("job1"); + CompactionJob job2 = createJobOnQueue("job2"); + + // When + RecordsProcessed job1RecordsProcessed = new RecordsProcessed(10L, 10L); + RecordsProcessed job2RecordsProcessed = new RecordsProcessed(5L, 5L); + runTask("test-task-1", processJobs( + jobSucceeds(job1RecordsProcessed), + jobSucceeds(job2RecordsProcessed)), + times::poll); + + // Then + RecordsProcessedSummary job1Summary = new RecordsProcessedSummary(job1RecordsProcessed, + Instant.parse("2024-02-22T13:50:01Z"), + Instant.parse("2024-02-22T13:50:02Z")); + RecordsProcessedSummary job2Summary = new RecordsProcessedSummary(job2RecordsProcessed, + Instant.parse("2024-02-22T13:50:03Z"), + Instant.parse("2024-02-22T13:50:04Z")); + assertThat(taskStore.getAllTasks()).containsExactly( + finishedCompactionTask("test-task-1", + Instant.parse("2024-02-22T13:50:00Z"), + Instant.parse("2024-02-22T13:50:05Z"), + job1Summary, job2Summary)); + assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).containsExactlyInAnyOrder( + jobCreated(job1, DEFAULT_CREATED_TIME, + finishedCompactionRun("test-task-1", job1Summary)), + jobCreated(job2, DEFAULT_CREATED_TIME, + finishedCompactionRun("test-task-1", job2Summary))); + } + + @Test + void shouldSaveTaskAndJobWhenOneJobFails() throws Exception { + // Given + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // Job started + Instant.parse("2024-02-22T13:50:05Z"))); // Finish + CompactionJob job = createJobOnQueue("job1"); + + // When + runTask("test-task-1", processJobs(jobFails()), times::poll); + + // Then + assertThat(taskStore.getAllTasks()).containsExactly( + finishedCompactionTask("test-task-1", + Instant.parse("2024-02-22T13:50:00Z"), + Instant.parse("2024-02-22T13:50:05Z"))); + assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).containsExactly( + jobCreated(job, DEFAULT_CREATED_TIME, + startedCompactionRun("test-task-1", Instant.parse("2024-02-22T13:50:01Z")))); + } + + @Test + void shouldSaveTaskWhenNoJobsFound() throws Exception { + // Given + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:05Z"))); // Finish + + // When + runTask("test-task-1", processNoJobs(), times::poll); + + // Then + assertThat(taskStore.getAllTasks()).containsExactly( + finishedCompactionTask("test-task-1", + Instant.parse("2024-02-22T13:50:00Z"), + Instant.parse("2024-02-22T13:50:05Z"))); + assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).isEmpty(); + } + + private CompactionTaskFinishedStatus.Builder withJobSummaries(RecordsProcessedSummary... summaries) { + CompactionTaskFinishedStatus.Builder taskFinishedBuilder = CompactionTaskFinishedStatus.builder(); + Stream.of(summaries).forEach(taskFinishedBuilder::addJobSummary); + return taskFinishedBuilder; + } + + private CompactionTaskStatus finishedCompactionTask(String taskId, Instant startTime, Instant finishTime, RecordsProcessedSummary... summaries) { + return CompactionTaskStatus.builder() + .startTime(startTime) + .taskId(taskId) + .finished(finishTime, withJobSummaries(summaries)) + .build(); + } + } +} diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTerminateTest.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTerminateTest.java new file mode 100644 index 0000000000..0165db67e6 --- /dev/null +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTerminateTest.java @@ -0,0 +1,218 @@ +/* + * 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.compaction.job.execution; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +import sleeper.compaction.job.CompactionJob; + +import java.time.Duration; +import java.time.Instant; +import java.util.LinkedList; +import java.util.List; +import java.util.Queue; + +import static org.assertj.core.api.Assertions.assertThat; +import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS; +import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES; +import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS; + +public class CompactionTaskTerminateTest extends CompactionTaskTestBase { + + @Nested + @DisplayName("Stop if idle for a specified period") + class StopAfterMaxIdleTime { + + @Test + void shouldTerminateIfNoJobsArePresentAfterRunningForIdleTime() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); + instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:03Z"))); // Finish + + // When + runTask(processNoJobs(), times::poll); + + // Then + assertThat(times).isEmpty(); + assertThat(sleeps).isEmpty(); + } + + @Test + void shouldTerminateIfNoJobsArePresentAfterRunningForIdleTimeWithTwoQueuePolls() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); + instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:02Z"), // First idle time check + Instant.parse("2024-02-22T13:50:04Z"))); // Second idle time check + finish + + // When + runTask(processNoJobs(), times::poll); + + // Then + assertThat(times).isEmpty(); + assertThat(sleeps).containsExactly(Duration.ofSeconds(2)); + } + + @Test + void shouldTerminateAfterRunningJobAndWaitingForIdleTime() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); + instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // Job started + Instant.parse("2024-02-22T13:50:02Z"), // Job completed + Instant.parse("2024-02-22T13:50:05Z"))); // Idle time check with empty queue and finish + CompactionJob job = createJobOnQueue("job1"); + + // When + runTask(jobsSucceed(1), times::poll); + + // Then + assertThat(times).isEmpty(); + assertThat(successfulJobs).containsExactly(job); + assertThat(failedJobs).isEmpty(); + assertThat(jobsOnQueue).isEmpty(); + assertThat(sleeps).isEmpty(); + } + + @Test + void shouldTerminateWhenMaxIdleTimeNotMetOnFirstCheckThenIdleAfterProcessingJob() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); + instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // First check + Instant.parse("2024-02-22T13:50:02Z"), // Job started + Instant.parse("2024-02-22T13:50:02Z"), // Job completed + Instant.parse("2024-02-22T13:50:06Z"))); // Second check + finish + CompactionJob job = createJob("job1"); + + // When + runTask( + pollQueue( + receiveNoJobAnd(() -> send(job)), + receiveJob(), + receiveNoJob()), + processJobs(jobSucceeds()), + times::poll); + + // Then + assertThat(times).isEmpty(); + assertThat(successfulJobs).containsExactly(job); + assertThat(failedJobs).isEmpty(); + assertThat(jobsOnQueue).isEmpty(); + assertThat(sleeps).containsExactly(Duration.ofSeconds(2)); + } + + @Test + void shouldTerminateWhenMaxIdleTimeNotMetOnFirstCheckThenNotMetAfterProcessingJob() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); + instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:01Z"), // First check + Instant.parse("2024-02-22T13:50:02Z"), // Job started + Instant.parse("2024-02-22T13:50:03Z"), // Job completed + Instant.parse("2024-02-22T13:50:04Z"), // Second check + Instant.parse("2024-02-22T13:50:06Z"))); // Third check + finish + CompactionJob job = createJob("job1"); + + // When + runTask( + pollQueue( + receiveNoJobAnd(() -> send(job)), + receiveJob(), + receiveNoJob(), + receiveNoJob()), + processJobs(jobSucceeds()), + times::poll); + + // Then + assertThat(times).isEmpty(); + assertThat(successfulJobs).containsExactly(job); + assertThat(failedJobs).isEmpty(); + assertThat(jobsOnQueue).isEmpty(); + assertThat(sleeps).containsExactly(Duration.ofSeconds(2), Duration.ofSeconds(2)); + } + + @Test + void shouldNotDelayRetryIfSetToZero() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); + instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 0); + Queue times = new LinkedList<>(List.of( + Instant.parse("2024-02-22T13:50:00Z"), // Start + Instant.parse("2024-02-22T13:50:02Z"), // First idle time check + Instant.parse("2024-02-22T13:50:04Z"))); // Second idle time check + finish + + // When + runTask(processNoJobs(), times::poll); + + // Then + assertThat(times).isEmpty(); + assertThat(sleeps).isEmpty(); + } + } + + @Nested + @DisplayName("Stop if failed too many times consecutively") + class StopAfterConsecutiveFailures { + @Test + void shouldStopEarlyIfMaxConsecutiveFailuresMet() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES, 2); + CompactionJob job1 = createJobOnQueue("job1"); + CompactionJob job2 = createJobOnQueue("job2"); + CompactionJob job3 = createJobOnQueue("job3"); + + // When + runTask(processJobs(jobFails(), jobFails(), jobSucceeds())); + + // Then + assertThat(successfulJobs).isEmpty(); + assertThat(failedJobs).containsExactly(job1, job2); + assertThat(jobsOnQueue).containsExactly(job3); + } + + @Test + void shouldResetConsecutiveFailureCountIfJobProcessedSuccessfully() throws Exception { + // Given + instanceProperties.setNumber(COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES, 2); + CompactionJob job1 = createJobOnQueue("job1"); + CompactionJob job2 = createJobOnQueue("job2"); + CompactionJob job3 = createJobOnQueue("job3"); + CompactionJob job4 = createJobOnQueue("job4"); + + // When + runTask(processJobs(jobFails(), jobSucceeds(), jobFails(), jobSucceeds())); + + // Then + assertThat(successfulJobs).containsExactly(job2, job4); + assertThat(failedJobs).containsExactly(job1, job3); + assertThat(jobsOnQueue).isEmpty(); + } + } +} diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTest.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTest.java index 9e3569593f..569683587a 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTest.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTest.java @@ -16,64 +16,15 @@ package sleeper.compaction.job.execution; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import sleeper.compaction.job.CompactionJob; -import sleeper.compaction.job.CompactionRunner; -import sleeper.compaction.job.execution.CompactionTask.MessageHandle; -import sleeper.compaction.job.execution.CompactionTask.MessageReceiver; -import sleeper.compaction.task.CompactionTaskFinishedStatus; -import sleeper.compaction.task.CompactionTaskStatus; -import sleeper.compaction.task.CompactionTaskStatusStore; -import sleeper.compaction.testutils.InMemoryCompactionJobStatusStore; -import sleeper.compaction.testutils.InMemoryCompactionTaskStatusStore; -import sleeper.configuration.properties.PropertiesReloader; -import sleeper.configuration.properties.instance.InstanceProperties; -import sleeper.core.record.process.RecordsProcessed; -import sleeper.core.record.process.RecordsProcessedSummary; - -import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; -import java.util.Iterator; -import java.util.LinkedList; -import java.util.List; -import java.util.Optional; -import java.util.Queue; -import java.util.UUID; -import java.util.function.Supplier; -import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; -import static sleeper.compaction.job.CompactionJobStatusTestData.finishedCompactionRun; -import static sleeper.compaction.job.CompactionJobStatusTestData.jobCreated; -import static sleeper.compaction.job.CompactionJobStatusTestData.startedCompactionRun; -import static sleeper.configuration.properties.InstancePropertiesTestHelper.createTestInstanceProperties; -import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS; -import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES; -import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS; - -public class CompactionTaskTest { - private static final String DEFAULT_TABLE_ID = "test-table-id"; - private static final String DEFAULT_TASK_ID = "test-task-id"; - private static final Instant DEFAULT_CREATED_TIME = Instant.parse("2024-03-04T10:50:00Z"); - - private final InstanceProperties instanceProperties = createTestInstanceProperties(); - private final Queue jobsOnQueue = new LinkedList<>(); - private final List successfulJobs = new ArrayList<>(); - private final List failedJobs = new ArrayList<>(); - private final InMemoryCompactionJobStatusStore jobStore = new InMemoryCompactionJobStatusStore(); - private final CompactionTaskStatusStore taskStore = new InMemoryCompactionTaskStatusStore(); - private final List sleeps = new ArrayList<>(); - @BeforeEach - void setUp() { - instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 0); - instanceProperties.setNumber(COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES, 10); - } +public class CompactionTaskTest extends CompactionTaskTestBase { @Nested @DisplayName("Process jobs") @@ -122,493 +73,4 @@ void shouldProcessTwoJobsFromQueueThenTerminate() throws Exception { assertThat(jobsOnQueue).isEmpty(); } } - - @Nested - @DisplayName("Stop if idle for a specified period") - class StopAfterMaxIdleTime { - - @Test - void shouldTerminateIfNoJobsArePresentAfterRunningForIdleTime() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); - instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:03Z"))); // Finish - - // When - runTask(processNoJobs(), times::poll); - - // Then - assertThat(times).isEmpty(); - assertThat(sleeps).isEmpty(); - } - - @Test - void shouldTerminateIfNoJobsArePresentAfterRunningForIdleTimeWithTwoQueuePolls() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); - instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:02Z"), // First idle time check - Instant.parse("2024-02-22T13:50:04Z"))); // Second idle time check + finish - - // When - runTask(processNoJobs(), times::poll); - - // Then - assertThat(times).isEmpty(); - assertThat(sleeps).containsExactly(Duration.ofSeconds(2)); - } - - @Test - void shouldTerminateAfterRunningJobAndWaitingForIdleTime() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); - instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:01Z"), // Job started - Instant.parse("2024-02-22T13:50:02Z"), // Job completed - Instant.parse("2024-02-22T13:50:05Z"))); // Idle time check with empty queue and finish - CompactionJob job = createJobOnQueue("job1"); - - // When - runTask(jobsSucceed(1), times::poll); - - // Then - assertThat(times).isEmpty(); - assertThat(successfulJobs).containsExactly(job); - assertThat(failedJobs).isEmpty(); - assertThat(jobsOnQueue).isEmpty(); - assertThat(sleeps).isEmpty(); - } - - @Test - void shouldTerminateWhenMaxIdleTimeNotMetOnFirstCheckThenIdleAfterProcessingJob() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); - instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:01Z"), // First check - Instant.parse("2024-02-22T13:50:02Z"), // Job started - Instant.parse("2024-02-22T13:50:02Z"), // Job completed - Instant.parse("2024-02-22T13:50:06Z"))); // Second check + finish - CompactionJob job = createJob("job1"); - - // When - runTask( - pollQueue( - receiveNoJobAnd(() -> send(job)), - receiveJob(), - receiveNoJob()), - processJobs(jobSucceeds()), - times::poll); - - // Then - assertThat(times).isEmpty(); - assertThat(successfulJobs).containsExactly(job); - assertThat(failedJobs).isEmpty(); - assertThat(jobsOnQueue).isEmpty(); - assertThat(sleeps).containsExactly(Duration.ofSeconds(2)); - } - - @Test - void shouldTerminateWhenMaxIdleTimeNotMetOnFirstCheckThenNotMetAfterProcessingJob() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); - instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 2); - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:01Z"), // First check - Instant.parse("2024-02-22T13:50:02Z"), // Job started - Instant.parse("2024-02-22T13:50:03Z"), // Job completed - Instant.parse("2024-02-22T13:50:04Z"), // Second check - Instant.parse("2024-02-22T13:50:06Z"))); // Third check + finish - CompactionJob job = createJob("job1"); - - // When - runTask( - pollQueue( - receiveNoJobAnd(() -> send(job)), - receiveJob(), - receiveNoJob(), - receiveNoJob()), - processJobs(jobSucceeds()), - times::poll); - - // Then - assertThat(times).isEmpty(); - assertThat(successfulJobs).containsExactly(job); - assertThat(failedJobs).isEmpty(); - assertThat(jobsOnQueue).isEmpty(); - assertThat(sleeps).containsExactly(Duration.ofSeconds(2), Duration.ofSeconds(2)); - } - - @Test - void shouldNotDelayRetryIfSetToZero() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 3); - instanceProperties.setNumber(COMPACTION_TASK_DELAY_BEFORE_RETRY_IN_SECONDS, 0); - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:02Z"), // First idle time check - Instant.parse("2024-02-22T13:50:04Z"))); // Second idle time check + finish - - // When - runTask(processNoJobs(), times::poll); - - // Then - assertThat(times).isEmpty(); - assertThat(sleeps).isEmpty(); - } - } - - @Nested - @DisplayName("Stop if failed too many times consecutively") - class StopAfterConsecutiveFailures { - @Test - void shouldStopEarlyIfMaxConsecutiveFailuresMet() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES, 2); - CompactionJob job1 = createJobOnQueue("job1"); - CompactionJob job2 = createJobOnQueue("job2"); - CompactionJob job3 = createJobOnQueue("job3"); - - // When - runTask(processJobs(jobFails(), jobFails(), jobSucceeds())); - - // Then - assertThat(successfulJobs).isEmpty(); - assertThat(failedJobs).containsExactly(job1, job2); - assertThat(jobsOnQueue).containsExactly(job3); - } - - @Test - void shouldResetConsecutiveFailureCountIfJobProcessedSuccessfully() throws Exception { - // Given - instanceProperties.setNumber(COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES, 2); - CompactionJob job1 = createJobOnQueue("job1"); - CompactionJob job2 = createJobOnQueue("job2"); - CompactionJob job3 = createJobOnQueue("job3"); - CompactionJob job4 = createJobOnQueue("job4"); - - // When - runTask(processJobs(jobFails(), jobSucceeds(), jobFails(), jobSucceeds())); - - // Then - assertThat(successfulJobs).containsExactly(job2, job4); - assertThat(failedJobs).containsExactly(job1, job3); - assertThat(jobsOnQueue).isEmpty(); - } - } - - @Nested - @DisplayName("Update status stores") - class UpdateStatusStores { - @Test - void shouldSaveTaskAndJobWhenOneJobSucceeds() throws Exception { - // Given - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:01Z"), // Job started - Instant.parse("2024-02-22T13:50:02Z"), // Job completed - Instant.parse("2024-02-22T13:50:05Z"))); // Finish - CompactionJob job = createJobOnQueue("job1"); - - // When - RecordsProcessed recordsProcessed = new RecordsProcessed(10L, 10L); - runTask("test-task-1", processJobs( - jobSucceeds(recordsProcessed)), - times::poll); - - // Then - RecordsProcessedSummary jobSummary = new RecordsProcessedSummary(recordsProcessed, - Instant.parse("2024-02-22T13:50:01Z"), - Instant.parse("2024-02-22T13:50:02Z")); - assertThat(taskStore.getAllTasks()).containsExactly( - finishedCompactionTask("test-task-1", - Instant.parse("2024-02-22T13:50:00Z"), - Instant.parse("2024-02-22T13:50:05Z"), - jobSummary)); - assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).containsExactly( - jobCreated(job, DEFAULT_CREATED_TIME, - finishedCompactionRun("test-task-1", jobSummary))); - } - - @Test - void shouldSaveTaskAndJobsWhenMultipleJobsSucceed() throws Exception { - // Given - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:01Z"), // Job 1 started - Instant.parse("2024-02-22T13:50:02Z"), // Job 1 completed - Instant.parse("2024-02-22T13:50:03Z"), // Job 2 started - Instant.parse("2024-02-22T13:50:04Z"), // Job 2 completed - Instant.parse("2024-02-22T13:50:05Z"))); // Finish - CompactionJob job1 = createJobOnQueue("job1"); - CompactionJob job2 = createJobOnQueue("job2"); - - // When - RecordsProcessed job1RecordsProcessed = new RecordsProcessed(10L, 10L); - RecordsProcessed job2RecordsProcessed = new RecordsProcessed(5L, 5L); - runTask("test-task-1", processJobs( - jobSucceeds(job1RecordsProcessed), - jobSucceeds(job2RecordsProcessed)), - times::poll); - - // Then - RecordsProcessedSummary job1Summary = new RecordsProcessedSummary(job1RecordsProcessed, - Instant.parse("2024-02-22T13:50:01Z"), - Instant.parse("2024-02-22T13:50:02Z")); - RecordsProcessedSummary job2Summary = new RecordsProcessedSummary(job2RecordsProcessed, - Instant.parse("2024-02-22T13:50:03Z"), - Instant.parse("2024-02-22T13:50:04Z")); - assertThat(taskStore.getAllTasks()).containsExactly( - finishedCompactionTask("test-task-1", - Instant.parse("2024-02-22T13:50:00Z"), - Instant.parse("2024-02-22T13:50:05Z"), - job1Summary, job2Summary)); - assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).containsExactlyInAnyOrder( - jobCreated(job1, DEFAULT_CREATED_TIME, - finishedCompactionRun("test-task-1", job1Summary)), - jobCreated(job2, DEFAULT_CREATED_TIME, - finishedCompactionRun("test-task-1", job2Summary))); - } - - @Test - void shouldSaveTaskAndJobWhenOneJobFails() throws Exception { - // Given - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:01Z"), // Job started - Instant.parse("2024-02-22T13:50:05Z"))); // Finish - CompactionJob job = createJobOnQueue("job1"); - - // When - runTask("test-task-1", processJobs(jobFails()), times::poll); - - // Then - assertThat(taskStore.getAllTasks()).containsExactly( - finishedCompactionTask("test-task-1", - Instant.parse("2024-02-22T13:50:00Z"), - Instant.parse("2024-02-22T13:50:05Z"))); - assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).containsExactly( - jobCreated(job, DEFAULT_CREATED_TIME, - startedCompactionRun("test-task-1", Instant.parse("2024-02-22T13:50:01Z")))); - } - - @Test - void shouldSaveTaskWhenNoJobsFound() throws Exception { - // Given - Queue times = new LinkedList<>(List.of( - Instant.parse("2024-02-22T13:50:00Z"), // Start - Instant.parse("2024-02-22T13:50:05Z"))); // Finish - - // When - runTask("test-task-1", processNoJobs(), times::poll); - - // Then - assertThat(taskStore.getAllTasks()).containsExactly( - finishedCompactionTask("test-task-1", - Instant.parse("2024-02-22T13:50:00Z"), - Instant.parse("2024-02-22T13:50:05Z"))); - assertThat(jobStore.getAllJobs(DEFAULT_TABLE_ID)).isEmpty(); - } - - private CompactionTaskFinishedStatus.Builder withJobSummaries(RecordsProcessedSummary... summaries) { - CompactionTaskFinishedStatus.Builder taskFinishedBuilder = CompactionTaskFinishedStatus.builder(); - Stream.of(summaries).forEach(taskFinishedBuilder::addJobSummary); - return taskFinishedBuilder; - } - - private CompactionTaskStatus finishedCompactionTask(String taskId, Instant startTime, Instant finishTime, RecordsProcessedSummary... summaries) { - return CompactionTaskStatus.builder() - .startTime(startTime) - .taskId(taskId) - .finished(finishTime, withJobSummaries(summaries)) - .build(); - } - } - - private void runTask(CompactionRunner compactor) throws Exception { - runTask(compactor, Instant::now); - } - - private void runTask(CompactionRunner compactor, Supplier timeSupplier) throws Exception { - runTask(pollQueue(), compactor, timeSupplier, DEFAULT_TASK_ID); - } - - private void runTask(String taskId, CompactionRunner compactor, Supplier timeSupplier) throws Exception { - runTask(pollQueue(), compactor, timeSupplier, taskId); - } - - private void runTask( - MessageReceiver messageReceiver, - CompactionRunner compactor, - Supplier timeSupplier) throws Exception { - runTask(messageReceiver, compactor, timeSupplier, DEFAULT_TASK_ID); - } - - private void runTask( - MessageReceiver messageReceiver, - CompactionRunner compactor, - Supplier timeSupplier, - String taskId) throws Exception { - CompactionAlgorithmSelector identity = (job) -> compactor; - new CompactionTask(instanceProperties, PropertiesReloader.neverReload(), - messageReceiver, identity, jobStore, taskStore, taskId, timeSupplier, sleeps::add) - .run(); - } - - private CompactionJob createJobOnQueue(String jobId) { - CompactionJob job = createJob(jobId); - jobsOnQueue.add(job); - jobStore.jobCreated(job, DEFAULT_CREATED_TIME); - return job; - } - - private CompactionJob createJob(String jobId) { - return CompactionJob.builder() - .tableId(DEFAULT_TABLE_ID) - .jobId(jobId) - .partitionId("root") - .inputFiles(List.of(UUID.randomUUID().toString())) - .outputFile(UUID.randomUUID().toString()).build(); - } - - private void send(CompactionJob job) { - jobsOnQueue.add(job); - } - - private MessageReceiver pollQueue() { - return () -> { - CompactionJob job = jobsOnQueue.poll(); - if (job != null) { - return Optional.of(new FakeMessageHandle(job)); - } else { - return Optional.empty(); - } - }; - } - - private MessageReceiver pollQueue(MessageReceiver... actions) { - Iterator getAction = List.of(actions).iterator(); - return () -> { - if (getAction.hasNext()) { - return getAction.next().receiveMessage(); - } else { - throw new IllegalStateException("Unexpected queue poll"); - } - }; - } - - private MessageReceiver receiveJob() { - return () -> { - if (jobsOnQueue.isEmpty()) { - throw new IllegalStateException("Expected job on queue"); - } - return Optional.of(new FakeMessageHandle(jobsOnQueue.poll())); - }; - } - - private MessageReceiver receiveNoJob() { - return () -> { - if (!jobsOnQueue.isEmpty()) { - throw new IllegalStateException("Expected no jobs on queue"); - } - return Optional.empty(); - }; - } - - private MessageReceiver receiveNoJobAnd(Runnable action) { - return () -> { - if (!jobsOnQueue.isEmpty()) { - throw new IllegalStateException("Expected no jobs on queue"); - } - action.run(); - return Optional.empty(); - }; - } - - private CompactionRunner jobsSucceed(int numJobs) { - return processJobs(Stream.generate(() -> jobSucceeds()) - .limit(numJobs) - .toArray(ProcessJob[]::new)); - } - - private ProcessJob jobSucceeds(RecordsProcessed summary) { - return new ProcessJob(true, summary); - } - - private ProcessJob jobSucceeds() { - return new ProcessJob(true, 10L); - } - - private ProcessJob jobFails() { - return new ProcessJob(false, 0L); - } - - private CompactionRunner processNoJobs() { - return processJobs(); - } - - private CompactionRunner processJobs(ProcessJob... actions) { - Iterator getAction = List.of(actions).iterator(); - return job -> { - if (getAction.hasNext()) { - ProcessJob action = getAction.next(); - action.run(job); - return action.summary; - } else { - throw new IllegalStateException("Unexpected job: " + job); - } - }; - } - - private class ProcessJob { - private final boolean succeed; - private final RecordsProcessed summary; - - ProcessJob(boolean succeed, long records) { - this(succeed, new RecordsProcessed(records, records)); - } - - ProcessJob(boolean succeed, RecordsProcessed summary) { - this.succeed = succeed; - this.summary = summary; - } - - public void run(CompactionJob job) throws Exception { - if (succeed) { - successfulJobs.add(job); - } else { - throw new Exception("Failed to process job"); - } - } - } - - private class FakeMessageHandle implements MessageHandle { - private final CompactionJob job; - - FakeMessageHandle(CompactionJob job) { - this.job = job; - } - - public CompactionJob getJob() { - return job; - } - - public void close() { - } - - public void completed() { - } - - public void failed() { - failedJobs.add(job); - } - } } diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTestBase.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTestBase.java new file mode 100644 index 0000000000..98913d60f3 --- /dev/null +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/CompactionTaskTestBase.java @@ -0,0 +1,312 @@ +/* + * 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.compaction.job.execution; + +import org.junit.jupiter.api.BeforeEach; + +import sleeper.compaction.job.CompactionJob; +import sleeper.compaction.job.commit.CompactionJobCommitRequest; +import sleeper.compaction.job.commit.CompactionJobCommitter; +import sleeper.compaction.job.execution.CompactionTask.CompactionRunner; +import sleeper.compaction.job.execution.CompactionTask.MessageHandle; +import sleeper.compaction.job.execution.CompactionTask.MessageReceiver; +import sleeper.compaction.task.CompactionTaskStatusStore; +import sleeper.compaction.testutils.InMemoryCompactionJobStatusStore; +import sleeper.compaction.testutils.InMemoryCompactionTaskStatusStore; +import sleeper.configuration.properties.PropertiesReloader; +import sleeper.configuration.properties.instance.InstanceProperties; +import sleeper.configuration.properties.table.FixedTablePropertiesProvider; +import sleeper.configuration.properties.table.TableProperties; +import sleeper.configuration.properties.table.TablePropertiesProvider; +import sleeper.core.record.process.RecordsProcessed; +import sleeper.core.schema.Schema; +import sleeper.core.statestore.FileReferenceFactory; +import sleeper.core.statestore.StateStore; +import sleeper.statestore.FixedStateStoreProvider; +import sleeper.statestore.StateStoreProvider; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.LinkedList; +import java.util.List; +import java.util.Optional; +import java.util.Queue; +import java.util.UUID; +import java.util.function.Supplier; +import java.util.stream.Stream; + +import static sleeper.configuration.properties.InstancePropertiesTestHelper.createTestInstanceProperties; +import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES; +import static sleeper.configuration.properties.instance.CompactionProperty.COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS; +import static sleeper.configuration.properties.table.TablePropertiesTestHelper.createTestTableProperties; +import static sleeper.configuration.properties.table.TableProperty.TABLE_ID; +import static sleeper.configuration.properties.table.TableProperty.TABLE_NAME; +import static sleeper.core.schema.SchemaTestHelper.schemaWithKey; +import static sleeper.core.statestore.AssignJobIdRequest.assignJobOnPartitionToFiles; +import static sleeper.core.statestore.inmemory.StateStoreTestHelper.inMemoryStateStoreWithSinglePartition; + +public class CompactionTaskTestBase { + protected static final String DEFAULT_TABLE_ID = "test-table-id"; + protected static final String DEFAULT_TABLE_NAME = "test-table-name"; + protected static final String DEFAULT_TASK_ID = "test-task-id"; + protected static final Instant DEFAULT_CREATED_TIME = Instant.parse("2024-03-04T10:50:00Z"); + + protected final InstanceProperties instanceProperties = createTestInstanceProperties(); + protected final Schema schema = schemaWithKey("key"); + protected final TableProperties tableProperties = createTable(DEFAULT_TABLE_ID, DEFAULT_TABLE_NAME); + protected final StateStore stateStore = inMemoryStateStoreWithSinglePartition(schema); + protected final FileReferenceFactory factory = FileReferenceFactory.from(stateStore); + protected final Queue jobsOnQueue = new LinkedList<>(); + protected final List successfulJobs = new ArrayList<>(); + protected final List failedJobs = new ArrayList<>(); + protected final InMemoryCompactionJobStatusStore jobStore = new InMemoryCompactionJobStatusStore(); + protected final CompactionTaskStatusStore taskStore = new InMemoryCompactionTaskStatusStore(); + protected final List sleeps = new ArrayList<>(); + protected final List commitRequestsOnQueue = new ArrayList<>(); + + @BeforeEach + void setUp() { + instanceProperties.setNumber(COMPACTION_TASK_MAX_IDLE_TIME_IN_SECONDS, 0); + instanceProperties.setNumber(COMPACTION_TASK_MAX_CONSECUTIVE_FAILURES, 10); + } + + protected TableProperties createTable(String tableId, String tableName) { + TableProperties tableProperties = createTestTableProperties(instanceProperties, schema); + tableProperties.set(TABLE_ID, tableId); + tableProperties.set(TABLE_NAME, tableName); + return tableProperties; + } + + protected void runTask(CompactionRunner compactor) throws Exception { + runTask(compactor, Instant::now); + } + + protected void runTask(CompactionRunner compactor, Supplier timeSupplier) throws Exception { + runTask(pollQueue(), compactor, timeSupplier, DEFAULT_TASK_ID); + } + + protected void runTask(CompactionRunner compactor, Supplier timeSupplier, + TablePropertiesProvider tablePropertiesProvider, StateStoreProvider stateStoreProvider) throws Exception { + runTask(pollQueue(), compactor, timeSupplier, DEFAULT_TASK_ID, tablePropertiesProvider, stateStoreProvider); + } + + protected void runTask(String taskId, CompactionRunner compactor, Supplier timeSupplier) throws Exception { + runTask(pollQueue(), compactor, timeSupplier, taskId); + } + + protected void runTask( + MessageReceiver messageReceiver, + CompactionRunner compactor, + Supplier timeSupplier) throws Exception { + runTask(messageReceiver, compactor, timeSupplier, DEFAULT_TASK_ID); + } + + private void runTask( + MessageReceiver messageReceiver, + CompactionRunner compactor, + Supplier timeSupplier, + String taskId) throws Exception { + runTask(messageReceiver, compactor, timeSupplier, taskId, + new FixedTablePropertiesProvider(tableProperties), + new FixedStateStoreProvider(tableProperties, stateStore)); + } + + private void runTask( + MessageReceiver messageReceiver, + CompactionRunner compactor, + Supplier timeSupplier, + String taskId, + TablePropertiesProvider tablePropertiesProvider, + StateStoreProvider stateStoreProvider) throws Exception { + CompactionJobCommitHandler commitHandler = new CompactionJobCommitHandler( + tablePropertiesProvider, + new CompactionJobCommitter(jobStore, tableId -> stateStoreProvider.getStateStore(tablePropertiesProvider.getById(tableId))), + commitRequestsOnQueue::add); + new CompactionTask(instanceProperties, + PropertiesReloader.neverReload(), messageReceiver, compactor, + commitHandler, jobStore, taskStore, taskId, timeSupplier, sleeps::add) + .run(); + } + + protected CompactionJob createJobOnQueue(String jobId) throws Exception { + return createJobOnQueue(jobId, tableProperties, stateStore); + } + + protected CompactionJob createJobOnQueue(String jobId, TableProperties tableProperties, StateStore stateStore) throws Exception { + CompactionJob job = createJob(jobId, tableProperties, stateStore); + jobsOnQueue.add(job); + jobStore.jobCreated(job, DEFAULT_CREATED_TIME); + return job; + } + + protected CompactionJob createJob(String jobId) throws Exception { + return createJob(jobId, tableProperties, stateStore); + } + + protected CompactionJob createJob(String jobId, TableProperties tableProperties, StateStore stateStore) throws Exception { + String inputFile = UUID.randomUUID().toString(); + CompactionJob job = CompactionJob.builder() + .tableId(tableProperties.get(TABLE_ID)) + .jobId(jobId) + .partitionId("root") + .inputFiles(List.of(inputFile)) + .outputFile(UUID.randomUUID().toString()).build(); + assignFilesToJob(job, stateStore); + return job; + } + + protected void assignFilesToJob(CompactionJob job, StateStore stateStore) throws Exception { + for (String inputFile : job.getInputFiles()) { + stateStore.addFile(factory.rootFile(inputFile, 123L)); + } + stateStore.assignJobIds(List.of(assignJobOnPartitionToFiles(job.getId(), job.getPartitionId(), job.getInputFiles()))); + } + + protected void send(CompactionJob job) { + jobsOnQueue.add(job); + } + + private MessageReceiver pollQueue() { + return () -> { + CompactionJob job = jobsOnQueue.poll(); + if (job != null) { + return Optional.of(new FakeMessageHandle(job)); + } else { + return Optional.empty(); + } + }; + } + + protected MessageReceiver pollQueue(MessageReceiver... actions) { + Iterator getAction = List.of(actions).iterator(); + return () -> { + if (getAction.hasNext()) { + return getAction.next().receiveMessage(); + } else { + throw new IllegalStateException("Unexpected queue poll"); + } + }; + } + + protected MessageReceiver receiveJob() { + return () -> { + if (jobsOnQueue.isEmpty()) { + throw new IllegalStateException("Expected job on queue"); + } + return Optional.of(new FakeMessageHandle(jobsOnQueue.poll())); + }; + } + + protected MessageReceiver receiveNoJob() { + return () -> { + if (!jobsOnQueue.isEmpty()) { + throw new IllegalStateException("Expected no jobs on queue"); + } + return Optional.empty(); + }; + } + + protected MessageReceiver receiveNoJobAnd(Runnable action) { + return () -> { + if (!jobsOnQueue.isEmpty()) { + throw new IllegalStateException("Expected no jobs on queue"); + } + action.run(); + return Optional.empty(); + }; + } + + protected CompactionRunner jobsSucceed(int numJobs) { + return processJobs(Stream.generate(() -> jobSucceeds()) + .limit(numJobs) + .toArray(ProcessJob[]::new)); + } + + protected ProcessJob jobSucceeds(RecordsProcessed summary) { + return new ProcessJob(true, summary); + } + + protected ProcessJob jobSucceeds() { + return new ProcessJob(true, 10L); + } + + protected ProcessJob jobFails() { + return new ProcessJob(false, 0L); + } + + protected CompactionRunner processNoJobs() { + return processJobs(); + } + + protected CompactionRunner processJobs(ProcessJob... actions) { + Iterator getAction = List.of(actions).iterator(); + return job -> { + if (getAction.hasNext()) { + ProcessJob action = getAction.next(); + action.run(job); + return action.summary; + } else { + throw new IllegalStateException("Unexpected job: " + job); + } + }; + } + + protected class ProcessJob { + private final boolean succeed; + private final RecordsProcessed summary; + + ProcessJob(boolean succeed, long records) { + this(succeed, new RecordsProcessed(records, records)); + } + + ProcessJob(boolean succeed, RecordsProcessed summary) { + this.succeed = succeed; + this.summary = summary; + } + + public void run(CompactionJob job) throws Exception { + if (succeed) { + successfulJobs.add(job); + } else { + throw new Exception("Failed to process job"); + } + } + } + + protected class FakeMessageHandle implements MessageHandle { + private final CompactionJob job; + + FakeMessageHandle(CompactionJob job) { + this.job = job; + } + + public CompactionJob getJob() { + return job; + } + + public void close() { + } + + public void completed() { + } + + public void failed() { + failedJobs.add(job); + } + } +} diff --git a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/ECSCompactionTaskRunnerLocalStackIT.java b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/ECSCompactionTaskRunnerLocalStackIT.java index c806906baf..8481ed1426 100644 --- a/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/ECSCompactionTaskRunnerLocalStackIT.java +++ b/java/compaction/compaction-job-execution/src/test/java/sleeper/compaction/job/execution/ECSCompactionTaskRunnerLocalStackIT.java @@ -39,6 +39,7 @@ import sleeper.compaction.job.CompactionJob; import sleeper.compaction.job.CompactionJobSerDe; import sleeper.compaction.job.CompactionJobStatusStore; +import sleeper.compaction.job.commit.CompactionJobCommitter; import sleeper.compaction.status.store.job.CompactionJobStatusStoreFactory; import sleeper.compaction.status.store.job.DynamoDBCompactionJobStatusStoreCreator; import sleeper.compaction.status.store.task.CompactionTaskStatusStoreFactory; @@ -341,9 +342,14 @@ private CompactionTask createTask(String taskId, StateStoreProvider stateStorePr DefaultSelector compactSortedFiles = new DefaultSelector(instanceProperties, tablePropertiesProvider, stateStoreProvider, ObjectFactory.noUserJars()); - CompactionTask task = new CompactionTask(instanceProperties, PropertiesReloader.neverReload(), - new SqsCompactionQueueHandler(sqs, instanceProperties), - compactSortedFiles, jobStatusStore, taskStatusStore, taskId); + CompactionJobCommitHandler commitHandler = new CompactionJobCommitHandler(tablePropertiesProvider, + new CompactionJobCommitter(jobStatusStore, tableId -> stateStoreProvider.getStateStore(tablePropertiesProvider.getById(tableId))), + (request) -> { + // TODO send to SQS and test once infrastructure is deployed by CDK + }); + CompactionTask task = new CompactionTask(instanceProperties, + PropertiesReloader.neverReload(), new SqsCompactionQueueHandler(sqs, instanceProperties), compactSortedFiles, + commitHandler, jobStatusStore, taskStatusStore, taskId); return task; } diff --git a/java/compaction/pom.xml b/java/compaction/pom.xml index 95148c7de6..0995d42404 100644 --- a/java/compaction/pom.xml +++ b/java/compaction/pom.xml @@ -32,7 +32,7 @@ compaction-job-creation-lambda compaction-job-execution compaction-rust - compaction-job-completion-lambda + compaction-job-committer-lambda compaction-status-store compaction-task-creation diff --git a/java/configuration/src/main/java/sleeper/configuration/properties/instance/DefaultProperty.java b/java/configuration/src/main/java/sleeper/configuration/properties/instance/DefaultProperty.java index 3d6f7a6d18..583f76d4e2 100644 --- a/java/configuration/src/main/java/sleeper/configuration/properties/instance/DefaultProperty.java +++ b/java/configuration/src/main/java/sleeper/configuration/properties/instance/DefaultProperty.java @@ -178,6 +178,12 @@ public interface DefaultProperty { "is large.") .defaultValue("async") .propertyGroup(InstancePropertyGroup.DEFAULT).build(); + UserDefinedInstanceProperty DEFAULT_COMPACTION_JOB_COMMIT_ASYNC = Index.propertyBuilder("sleeper.default.compaction.job.commit.async") + .description("If true, compaction job commit requests will be sent to the compaction job committer lambda " + + "to be performed asynchronously. If false, compaction jobs will be committed synchronously by compaction tasks.") + .defaultValue("false") + .validationPredicate(Utils::isTrueOrFalse) + .propertyGroup(InstancePropertyGroup.DEFAULT).build(); static List getAll() { return Index.INSTANCE.getAll(); diff --git a/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java b/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java index 86b0f643f6..951138212b 100644 --- a/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java +++ b/java/configuration/src/main/java/sleeper/configuration/properties/table/TableProperty.java @@ -36,6 +36,7 @@ import static sleeper.configuration.properties.instance.CompactionProperty.DEFAULT_SIZERATIO_COMPACTION_STRATEGY_RATIO; import static sleeper.configuration.properties.instance.DefaultProperty.DEFAULT_BULK_IMPORT_MIN_LEAF_PARTITION_COUNT; import static sleeper.configuration.properties.instance.DefaultProperty.DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH; +import static sleeper.configuration.properties.instance.DefaultProperty.DEFAULT_COMPACTION_JOB_COMMIT_ASYNC; import static sleeper.configuration.properties.instance.DefaultProperty.DEFAULT_COMPRESSION_CODEC; import static sleeper.configuration.properties.instance.DefaultProperty.DEFAULT_DICTIONARY_ENCODING_FOR_ROW_KEY_FIELDS; import static sleeper.configuration.properties.instance.DefaultProperty.DEFAULT_DICTIONARY_ENCODING_FOR_SORT_KEY_FIELDS; @@ -218,6 +219,12 @@ public interface TableProperty extends SleeperProperty { "assign job IDs to the input files.") .propertyGroup(TablePropertyGroup.COMPACTION) .build(); + TableProperty COMPACTION_JOB_COMMIT_ASYNC = Index.propertyBuilder("sleeper.table.compaction.job.commit.async") + .defaultProperty(DEFAULT_COMPACTION_JOB_COMMIT_ASYNC) + .description("If true, compaction job commit requests will be sent to the compaction job committer lambda " + + "to be performed asynchronously. If false, compaction jobs will be committed synchronously by compaction tasks.") + .propertyGroup(TablePropertyGroup.COMPACTION) + .build(); TableProperty SIZE_RATIO_COMPACTION_STRATEGY_RATIO = Index.propertyBuilder("sleeper.table.compaction.strategy.sizeratio.ratio") .defaultProperty(DEFAULT_SIZERATIO_COMPACTION_STRATEGY_RATIO) .description("Used by the SizeRatioCompactionStrategy to decide if a group of files should be compacted.\n" + diff --git a/java/pom.xml b/java/pom.xml index 84a7ee2862..88d7f1364a 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -100,7 +100,7 @@ OkHttp, as a dependency of Spark, declares an old version of Kotlin with vulnerabilities. Managed from a conflict between 1.6.20 and 1.5.31. --> - 1.9.23 + 1.9.24