1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.fs.storage
10 
11 import org.apache.parquet.filter2.predicate.FilterPredicate
12 import org.geotools.api.feature.simple.SimpleFeatureType
13 import org.geotools.api.filter.Filter
14 import org.geotools.feature.simple.SimpleFeatureTypeBuilder
15 import org.locationtech.geomesa.filter.FilterHelper
16 import org.locationtech.geomesa.utils.geotools.Transform.{ExpressionTransform, PropertyTransform, RenameTransform, Transforms}
17 
18 package object parquet {
19 
20   /**
21     * Schema to read and schema to return.
22     *
23     * If we have to return a different feature than we read, we need to apply a secondary transform.
24     * Otherwise, we can just do the transform on read and skip the secondary transform
25     *
26     * @param read read schema, includes fields to filter on
27     * @param transform return schema, if different from read schema
28     */
29   case class ReadSchema(read: SimpleFeatureType, transform: Option[(String, SimpleFeatureType)])
30 
31   /**
32     * Filter to read files
33     *
34     * @param parquet parquet filter that we can push down to the file format
35     * @param residual residual geotools filter that we have to apply after read
36     */
37   case class ReadFilter(parquet: Option[FilterPredicate], residual: Option[Filter])
38 
39   object ReadSchema {
40 
41     import org.locationtech.geomesa.filter.RichTransform.RichTransform
42 
43     /**
44       * Calculates the read schema
45       *
46       * @param sft simple feature type
47       * @param filter query filter
48       * @param transform query transform
49       * @return
50       */
51     def apply(
52         sft: SimpleFeatureType,
53         filter: Option[Filter],
54         transform: Option[(String, SimpleFeatureType)]): ReadSchema = {
55       transform match {
56         case None => ReadSchema(sft, None)
57         case Some((tdefs, _)) =>
58           val definitions = Transforms(sft, tdefs)
59           val secondary = definitions.exists {
60             case _: PropertyTransform   => false
61             case _: RenameTransform     => false
62             case _: ExpressionTransform => true
63           }
64           val transformCols = definitions.flatMap(_.properties).distinct
65           val filterCols = filter match {
66             case None => Seq.empty
67             case Some(f) => FilterHelper.propertyNames(f, sft).filterNot(transformCols.contains)
68           }
69 
70           val projectedSft = {
71             val builder = new SimpleFeatureTypeBuilder()
72             builder.setName(sft.getName)
73             transformCols.foreach(a => builder.add(sft.getDescriptor(a)))
74             filterCols.foreach(a => builder.add(sft.getDescriptor(a)))
75             builder.buildFeatureType()
76           }
77           projectedSft.getUserData.putAll(sft.getUserData)
78 
79           ReadSchema(projectedSft, if (secondary || filterCols.nonEmpty) { transform } else { None })
80       }
81     }
82   }
83 
84   object ReadFilter {
85 
86     /**
87       * Create a read filter
88       *
89       * @param sft simple feature type
90       * @param filter query filter
91       * @return
92       */
93     def apply(sft: SimpleFeatureType, filter: Option[Filter]): ReadFilter = {
94       val (parquet, residual) = filter match {
95         case None | Some(Filter.INCLUDE) => (None, None)
96         case Some(f) => FilterConverter.convert(sft, f)
97       }
98       ReadFilter(parquet, residual)
99     }
100   }
101 }