6138 - Spark client

.github/config/chunks.yaml

@@ -4,6 +4,7 @@ chunks:
     workflow: chunk-clients-cdk.yaml
     modules:
       - clients
+      - spark
       - deployment/cdk
       - deployment/cdk-custom-resources
       - deployment/cdk-environment

.github/workflows/chunk-clients-cdk.yaml

@@ -11,6 +11,7 @@ on:
       - 'java/deployment/pom.xml'
       - 'java/system-test/pom.xml'
       - 'java/clients/**'
+      - 'java/spark/**'
       - 'java/deployment/cdk/**'
       - 'java/deployment/cdk-custom-resources/**'
       - 'java/deployment/cdk-environment/**'

java/pom.xml

@@ -46,6 +46,7 @@
         <module>statestore-lambda</module>
         <module>system-test</module>
         <module>trino</module>
+        <module>spark</module>
     </modules>
 
     <properties>

java/spark/pom.xml

@@ -0,0 +1,86 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Copyright 2022-2025 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.
+  -->
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>aws</artifactId>
+        <groupId>sleeper</groupId>
+        <version>0.35.0-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>spark</artifactId>
+
+    <dependencies>
+    <!-- Spark -->
+        <dependency>
+            <groupId>org.apache.spark</groupId>
+            <artifactId>spark-sql_${scala.version}</artifactId>
+            <scope>provided</scope>
+            <exclusions>
+                <exclusion>
+                    <groupId>org.apache.hadoop</groupId>
+                    <artifactId>*</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+        <!-- Sleeper dependencies -->
+        <dependency>
+            <groupId>sleeper</groupId>
+            <artifactId>configuration</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>sleeper</groupId>
+            <artifactId>query-datafusion</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>sleeper</groupId>
+            <artifactId>statestore</artifactId>
+            <version>${project.parent.version}</version>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-shade-plugin</artifactId>
+                <configuration>
+                    <relocations>
+                        <relocation>
+                            <pattern>com.google</pattern>
+                            <shadedPattern>sleeper.shaded.com.google</shadedPattern>
+                            <excludes>
+                                <exclude>com.google.common.util.concurrent.*</exclude>
+                            </excludes>
+                        </relocation>
+                        <relocation>
+                            <pattern>org.apache.parquet</pattern>
+                            <shadedPattern>sleeper.shaded.apache.parquet</shadedPattern>
+                        </relocation>
+                        <relocation>
+                            <pattern>shaded.parquet</pattern>
+                            <shadedPattern>sleeper.shaded.parquet.shaded</shadedPattern>
+                        </relocation>
+                    </relocations>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
+</project>

java/spark/src/main/java/sleeper/spark/CreateRegionFromFilter.java

@@ -0,0 +1,198 @@
+/*
+ * Copyright 2022-2025 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.spark;
+
+import org.apache.spark.sql.sources.EqualTo;
+import org.apache.spark.sql.sources.Filter;
+import org.apache.spark.sql.sources.GreaterThan;
+import org.apache.spark.sql.sources.GreaterThanOrEqual;
+import org.apache.spark.sql.sources.In;
+import org.apache.spark.sql.sources.LessThan;
+import org.apache.spark.sql.sources.LessThanOrEqual;
+import org.apache.spark.sql.sources.Or;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import sleeper.core.range.Range;
+import sleeper.core.range.Range.RangeFactory;
+import sleeper.core.range.Region;
+import sleeper.core.schema.Field;
+import sleeper.core.schema.Schema;
+import sleeper.core.schema.type.PrimitiveType;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Creates Regions from filters.
+ */
+public class CreateRegionFromFilter {
+    private static final Logger LOGGER = LoggerFactory.getLogger(CreateRegionFromFilter.class);
+
+    private CreateRegionFromFilter() {
+
+    }
+
+    /**
+     * Converts a filter into an optional list of regions corresponding to that filter.
+     *
+     * @param  filter the filter
+     * @param  schema the schema of the Sleeper table
+     * @return        an optional list of regions corresponding to the provided filter
+     */
+    public static Optional<List<Region>> createRegionsFromFilter(Filter filter, Schema schema) {
+        if (filter instanceof Or) {
+            Or or = (Or) filter;
+            Region leftRegion = createRegionFromSimpleFilter(or.left(), schema);
+            Region rightRegion = createRegionFromSimpleFilter(or.right(), schema);
+            return Optional.of(List.of(leftRegion, rightRegion));
+        } else if (filter instanceof In) {
+            In in = (In) filter;
+            List<Region> regions = new ArrayList<>();
+            Object[] values = in.values();
+            for (Object value : values) {
+                Filter equalFilter = new EqualTo(in.attribute(), value);
+                Region optional = createRegionFromSimpleFilter(equalFilter, schema);
+                regions.add(optional);
+            }
+            if (regions.isEmpty()) {
+                return Optional.empty();
+            } else {
+                return Optional.of(regions);
+            }
+        } else {
+            Region optionalRegion = createRegionFromSimpleFilter(filter, schema);
+            return Optional.of(List.of(optionalRegion));
+        }
+    }
+
+    /**
+     * A region corresponding to the provided filter.
+     *
+     * @param  filter the filter
+     * @param  schema the schema of the Sleeper table
+     * @return        a region corresponding to the provided filter
+     */
+    public static Region createRegionFromSimpleFilter(Filter filter, Schema schema) {
+        MutableRegion mutableRegion = new MutableRegion(schema);
+        updateRegionWithFilter(mutableRegion, filter, new HashSet<>(schema.getRowKeyFieldNames()));
+        Region region = mutableRegion.getRegion();
+        LOGGER.debug("Filter {} was converted to Region {}", filter, region);
+        return region;
+    }
+
+    /**
+     * An internal class used to allow a region to be updated as filters are added.
+     */
+    private static class MutableRegion {
+        private final Map<String, MutableRange> rowKeyToMutableRange;
+
+        MutableRegion(Schema schema) {
+            this.rowKeyToMutableRange = new HashMap<>();
+            RangeFactory rangeFactory = new RangeFactory(schema);
+            for (Field rowKeyField : schema.getRowKeyFields()) {
+                this.rowKeyToMutableRange.put(rowKeyField.getName(), new MutableRange(rowKeyField, rangeFactory));
+            }
+        }
+
+        MutableRange getMutableRange(String rowKeyFieldName) {
+            return rowKeyToMutableRange.get(rowKeyFieldName);
+        }
+
+        Region getRegion() {
+            return new Region(rowKeyToMutableRange.values().stream().map(mr -> mr.getRange()).collect(Collectors.toList()));
+        }
+    }
+
+    /**
+     * An internal class used to allow a range to be updated as filters are added.
+     */
+    private static class MutableRange {
+        private final String fieldName;
+        private final RangeFactory rangeFactory;
+        Object min;
+        boolean minInclusive = true;
+        Object max = null;
+        boolean maxInclusive = false;
+
+        MutableRange(Field field, RangeFactory rangeFactory) {
+            this.fieldName = field.getName();
+            this.rangeFactory = rangeFactory;
+            this.min = PrimitiveType.getMinimum(field.getType());
+        }
+
+        Range getRange() {
+            return rangeFactory.createRange(fieldName, min, minInclusive, max, maxInclusive);
+        }
+    }
+
+    private static void updateRegionWithFilter(MutableRegion mutableRegion, Filter filter, Set<String> rowKeyFieldNames) {
+        if (filter instanceof EqualTo) {
+            String fieldName = ((EqualTo) filter).attribute();
+            if (!rowKeyFieldNames.contains(fieldName)) {
+                return;
+            }
+            Object wantedKey = ((EqualTo) filter).value();
+            MutableRange mutableRange = mutableRegion.getMutableRange(fieldName);
+            mutableRange.minInclusive = true;
+            mutableRange.min = wantedKey;
+            mutableRange.maxInclusive = true;
+            mutableRange.max = wantedKey;
+        } else if (filter instanceof GreaterThan) {
+            String fieldName = ((GreaterThan) filter).attribute();
+            if (!rowKeyFieldNames.contains(fieldName)) {
+                return;
+            }
+            Object minimumKey = ((GreaterThan) filter).value();
+            MutableRange mutableRange = mutableRegion.getMutableRange(fieldName);
+            mutableRange.minInclusive = false;
+            mutableRange.min = minimumKey;
+        } else if (filter instanceof GreaterThanOrEqual) {
+            String fieldName = ((GreaterThanOrEqual) filter).attribute();
+            if (!rowKeyFieldNames.contains(fieldName)) {
+                return;
+            }
+            Object minimumKey = ((GreaterThanOrEqual) filter).value();
+            MutableRange mutableRange = mutableRegion.getMutableRange(fieldName);
+            mutableRange.minInclusive = true;
+            mutableRange.min = minimumKey;
+        } else if (filter instanceof LessThan) {
+            String fieldName = ((LessThan) filter).attribute();
+            if (!rowKeyFieldNames.contains(fieldName)) {
+                return;
+            }
+            Object maxiumKey = ((LessThan) filter).value();
+            MutableRange mutableRange = mutableRegion.getMutableRange(fieldName);
+            mutableRange.maxInclusive = false;
+            mutableRange.max = maxiumKey;
+        } else if (filter instanceof LessThanOrEqual) {
+            String fieldName = ((LessThanOrEqual) filter).attribute();
+            if (!rowKeyFieldNames.contains(fieldName)) {
+                return;
+            }
+            Object maxiumKey = ((LessThanOrEqual) filter).value();
+            MutableRange mutableRange = mutableRegion.getMutableRange(fieldName);
+            mutableRange.maxInclusive = true;
+            mutableRange.max = maxiumKey;
+        }
+    }
+}

java/spark/src/main/java/sleeper/spark/CreateRegionsFromPushedFilters.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2022-2025 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.spark;
+
+import org.apache.spark.sql.sources.Filter;
+
+import sleeper.core.range.Range.RangeFactory;
+import sleeper.core.range.Region;
+import sleeper.core.schema.Schema;
+import sleeper.spark.SplitPushedFiltersIntoSingleAndMultiRegionFilters.SingleAndMultiRegionFilters;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * Used to identify a list of the minimum Regions that the pushed filters correspond to.
+ */
+public class CreateRegionsFromPushedFilters {
+    private final Schema schema;
+    private final RangeFactory rangeFactory;
+
+    public CreateRegionsFromPushedFilters(Schema schema) {
+        this.schema = schema;
+        this.rangeFactory = new RangeFactory(schema);
+    }
+
+    /**
+     * Given the filters that have been pushed to Sleeper, identifies the minimum list of regions that need to
+     * be read to return results corresponding to those filters. Note that this is not a one-to-one map between
+     * the pushedFilters array and the list of regions.
+     *
+     * @param  pushedFilters the filters that have been pushed to Sleeper
+     * @return               a list of the smallest regions that correspond to the filters
+     */
+    public List<Region> getMinimumRegionCoveringPushedFilters(Filter[] pushedFilters) {
+        SplitPushedFiltersIntoSingleAndMultiRegionFilters split = new SplitPushedFiltersIntoSingleAndMultiRegionFilters();
+        SingleAndMultiRegionFilters singleAndMultiRegionFilters = split.splitPushedFilters(pushedFilters);
+        Region regionFromSingleRegionFilters = getRegionFromSingleRegionFilters(singleAndMultiRegionFilters.getSingleRegionFilters());
+        List<Region> regionsFromMultiRegionFilters = getRegionsFromMultiRegionFilters(singleAndMultiRegionFilters.getMultiRegionFilters());
+        if (!regionsFromMultiRegionFilters.isEmpty()) {
+            List<Region> regions = new ArrayList<>();
+            for (Region region : regionsFromMultiRegionFilters) {
+                Optional<Region> optionalIntersectedRegion = RegionIntersector.intersectRegions(region, regionFromSingleRegionFilters, rangeFactory, schema);
+                if (optionalIntersectedRegion.isPresent()) {
+                    regions.add(optionalIntersectedRegion.get());
+                }
+            }
+            return regions;
+        } else {
+            return List.of(regionFromSingleRegionFilters);
+        }
+    }
+
+    private Region getRegionFromSingleRegionFilters(List<Filter> singleRegionFilters) {
+        if (singleRegionFilters == null || singleRegionFilters.isEmpty()) {
+            return Region.coveringAllValuesOfAllRowKeys(schema);
+        }
+        Region intersectedRegion = CreateRegionFromFilter.createRegionFromSimpleFilter(singleRegionFilters.get(0), schema);
+        for (int i = 1; i < singleRegionFilters.size(); i++) {
+            Region region = CreateRegionFromFilter.createRegionFromSimpleFilter(singleRegionFilters.get(i), schema);
+            intersectedRegion = RegionIntersector.intersectRegions(intersectedRegion, region, new RangeFactory(schema), schema).get();
+        }
+        return intersectedRegion;
+    }
+
+    private List<Region> getRegionsFromMultiRegionFilters(List<Filter> multiRegionFilters) {
+        List<Region> regions = new ArrayList<>();
+        for (Filter filter : multiRegionFilters) {
+            regions.addAll(CreateRegionFromFilter.createRegionsFromFilter(filter, schema).get());
+        }
+        return regions;
+    }
+}