Initial work on type-safe filter predicates

Author: grouzen

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/Expr.scala

@@ -0,0 +1,92 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+import org.apache.parquet.filter2.predicate.FilterPredicate
+import org.apache.parquet.filter2.predicate.FilterApi
+import me.mnedokushev.zio.apache.parquet.core.Value
+import zio.prelude._
+
+sealed trait Expr[-A]
+
+object Expr {
+
+  final case class Column[A](path: String)(implicit val typeTag: TypeTag[A]) extends Expr[A] { self =>
+
+    def /[B: TypeTag](column: Column[B]): Column[B] =
+      Column(s"$path.${column.path}")
+
+    def >(value: A)(implicit ev: OperatorSupport.LessGreater[A]): Predicate[A] =
+      Predicate.Binary(self, value, Operator.Binary.GreaterThen())
+
+    def ===(value: A)(implicit ev: OperatorSupport.EqNotEq[A]): Predicate[A] =
+      Predicate.Binary(self, value, Operator.Binary.Eq())
+
+  }
+
+  sealed trait Predicate[A] extends Expr[A] { self =>
+
+    def not: Predicate[A] =
+      Predicate.Unary(self, Operator.Unary.Not[A]())
+
+    def and[B](other: Predicate[B]): Predicate[A] =
+      Predicate.Logical(self, other, Operator.Logical.And[A, B])
+
+    def or[B](other: Predicate[B]): Predicate[A] =
+      Predicate.Logical(self, other, Operator.Logical.Or[A, B])
+
+  }
+
+  object Predicate {
+
+    final case class Binary[A](column: Column[A], value: A, op: Operator.Binary[A]) extends Predicate[A]
+
+    final case class Unary[A](predicate: Predicate[A], op: Operator.Unary[A]) extends Predicate[A]
+
+    final case class Logical[A, B](left: Predicate[A], right: Predicate[B], op: Operator.Logical[A, B])
+        extends Predicate[A]
+
+  }
+
+  def compile[A](predicate: Predicate[A]): Either[String, FilterPredicate] =
+    predicate match {
+      case Predicate.Unary(predicate0, op)    =>
+        op match {
+          case Operator.Unary.Not() =>
+            compile(predicate0).map(FilterApi.not)
+        }
+      case Predicate.Logical(left, right, op) =>
+        (compile(left) <*> compile(right)).map { case (left0, right0) =>
+          op match {
+            case Operator.Logical.And() =>
+              FilterApi.and(left0, right0)
+            case Operator.Logical.Or()  =>
+              FilterApi.or(left0, right0)
+          }
+        }
+      case Predicate.Binary(column, value, op) =>
+        (column.typeTag, value) match {
+          case (TypeTag.String, v: String) =>
+            val column0 = FilterApi.binaryColumn(column.path)
+            val value0  = Value.string(v).value
+
+            op match {
+              case Operator.Binary.Eq()    =>
+                Right(FilterApi.eq(column0, value0))
+              case Operator.Binary.NotEq() =>
+                Right(FilterApi.notEq(column0, value0))
+              case _                       => ???
+            }
+          case (TypeTag.Int, v: Int)       =>
+            val column0 = FilterApi.intColumn(column.path)
+            val value0  = Int.box(Value.int(v).value)
+
+            op match {
+              case Operator.Binary.GreaterThen() =>
+                Right(FilterApi.gt(column0, value0))
+              case _                             => ???
+            }
+          case _                           => ???
+        }
+
+    }
+
+}

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/ExprAccessorBuilder.scala

@@ -0,0 +1,26 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+import zio.schema.AccessorBuilder
+import zio.schema.Schema
+
+class ExprAccessorBuilder extends AccessorBuilder {
+
+  override type Lens[F, S, A] = Expr.Column[A]
+
+  override type Prism[F, S, A] = Unit
+
+  override type Traversal[S, A] = Unit
+
+  override def makeLens[F, S, A](product: Schema.Record[S], term: Schema.Field[S, A]): Expr.Column[A] = {
+    implicit val typeTag: TypeTag[A] = TypeTag.deriveTypeTag(term.schema).get
+
+    Expr.Column(term.name.toString)
+  }
+
+  override def makePrism[F, S, A](sum: Schema.Enum[S], term: Schema.Case[S, A]): Prism[F, S, A] =
+    ()
+
+  override def makeTraversal[S, A](collection: Schema.Collection[S, A], element: Schema[A]): Traversal[S, A] =
+    ()
+
+}

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/Filter.scala

@@ -0,0 +1,27 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+import zio.schema._
+import me.mnedokushev.zio.apache.parquet.core._
+
+trait Filter[Columns0] {
+
+  type Columns
+
+  val columns: Columns0
+
+}
+
+object Filter {
+
+  def columns[A](implicit schema: Schema.Record[A]): schema.Accessors[Lens, Prism, Traversal] =
+    new Filter[schema.Accessors[Lens, Prism, Traversal]] {
+
+      val accessorBuilder = new ExprAccessorBuilder
+
+      override type Columns = schema.Accessors[accessorBuilder.Lens, accessorBuilder.Prism, accessorBuilder.Traversal]
+
+      override val columns: Columns = schema.makeAccessors(accessorBuilder)
+
+    }.columns
+
+}

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/Operator.scala

@@ -0,0 +1,32 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+sealed trait Operator
+
+object Operator {
+
+  sealed trait Binary[A] extends Operator
+
+  object Binary {
+    final case class Eq[A: OperatorSupport.EqNotEq]()              extends Binary[A]
+    final case class NotEq[A: OperatorSupport.EqNotEq]()           extends Binary[A]
+    final case class LessThen[A: OperatorSupport.LessGreater]()    extends Binary[A]
+    final case class LessEq[A: OperatorSupport.LessGreater]()      extends Binary[A]
+    final case class GreaterThen[A: OperatorSupport.LessGreater]() extends Binary[A]
+    final case class GreaterEq[A: OperatorSupport.LessGreater]()   extends Binary[A]
+
+  }
+
+  sealed trait Unary[A] extends Operator
+
+  object Unary {
+    final case class Not[A]() extends Unary[A]
+  }
+
+  sealed trait Logical[A, B] extends Operator
+
+  object Logical {
+    final case class And[A, B]() extends Logical[A, B]
+    final case class Or[A, B]()  extends Logical[A, B]
+  }
+
+}

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/OperatorSupport.scala

@@ -0,0 +1,30 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+sealed trait OperatorSupport[A] {
+  def typeTag: TypeTag[A]
+}
+
+object OperatorSupport {
+
+  abstract class LessGreater[A: TypeTag] extends OperatorSupport[A] {
+    override def typeTag: TypeTag[A] = implicitly[TypeTag[A]]
+  }
+
+  object LessGreater {
+    implicit case object Byte  extends LessGreater[Byte]
+    implicit case object Short extends LessGreater[Short]
+    implicit case object Int   extends LessGreater[Int]
+  }
+
+  abstract class EqNotEq[A: TypeTag] extends OperatorSupport[A] {
+    override def typeTag: TypeTag[A] = implicitly[TypeTag[A]]
+  }
+
+  object EqNotEq {
+    implicit case object String  extends EqNotEq[String]
+    implicit case object Boolean extends EqNotEq[Boolean]
+    implicit case object Byte    extends EqNotEq[Byte]
+    implicit case object Short   extends EqNotEq[Short]
+  }
+
+}

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/filter/TypeTag.scala

@@ -0,0 +1,36 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+import zio.schema.Schema
+import zio.schema.StandardType
+
+sealed trait TypeTag[+A]
+
+object TypeTag {
+
+  implicit case object String  extends TypeTag[String]
+  implicit case object Boolean extends TypeTag[Boolean]
+  implicit case object Byte    extends TypeTag[Byte]
+  implicit case object Short   extends TypeTag[Short]
+  implicit case object Int     extends TypeTag[Int]
+  implicit case object Long    extends TypeTag[Long]
+
+  def deriveTypeTag[A](schema: Schema[A]): Option[TypeTag[A]] =
+    schema match {
+      case s: Schema.Lazy[_]                 => deriveTypeTag(s.schema)
+      case s: Schema.Optional[_]             => deriveTypeTag(s.schema.asInstanceOf[Schema[A]])
+      case Schema.Primitive(standartType, _) => deriveTypeTag(standartType)
+      case _                                 => None
+    }
+
+  def deriveTypeTag[A](standartType: StandardType[A]): Option[TypeTag[A]] =
+    standartType match {
+      case StandardType.StringType => Some(TypeTag.String)
+      case StandardType.BoolType   => Some(TypeTag.Boolean)
+      case StandardType.ByteType   => Some(TypeTag.Byte)
+      case StandardType.ShortType  => Some(TypeTag.Short)
+      case StandardType.IntType    => Some(TypeTag.Int)
+      case StandardType.LongType   => Some(TypeTag.Long)
+      case _                       => None
+    }
+
+}

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/hadoop/ParquetReader.scala

@@ -11,12 +11,14 @@ import zio.schema.Schema
 import zio.stream._
 
 import java.io.IOException
+import me.mnedokushev.zio.apache.parquet.core.filter.Expr
+import org.apache.parquet.filter2.compat.FilterCompat
 
 trait ParquetReader[+A <: Product] {
 
   def readStream(path: Path): ZStream[Scope, Throwable, A]
 
-  def readChunk(path: Path): Task[Chunk[A]]
+  def readChunk[B](path: Path, filter: Option[Expr.Predicate[B]] = None): Task[Chunk[A]]
 
 }
 
@@ -39,10 +41,10 @@ final class ParquetReaderLive[A <: Product: Tag](
                 )
     } yield value
 
-  override def readChunk(path: Path): Task[Chunk[A]] =
+  override def readChunk[B](path: Path, filter: Option[Expr.Predicate[B]] = None): Task[Chunk[A]] =
     ZIO.scoped(
       for {
-        reader  <- build(path)
+        reader  <- build(path, filter)
         readNext = for {
                      value  <- ZIO.attemptBlockingIO(reader.read())
                      record <- if (value != null)
@@ -63,14 +65,25 @@ final class ParquetReaderLive[A <: Product: Tag](
       } yield builder.result()
     )
 
-  private def build(path: Path): ZIO[Scope, IOException, HadoopParquetReader[RecordValue]] =
+  private def build[B](
+    path: Path,
+    filter: Option[Expr.Predicate[B]] = None
+  ): ZIO[Scope, IOException, HadoopParquetReader[RecordValue]] =
     for {
-      inputFile <- path.toInputFileZIO(hadoopConf)
-      reader    <- ZIO.fromAutoCloseable(
-                     ZIO.attemptBlockingIO(
-                       new ParquetReader.Builder(inputFile, schema, schemaEncoder).withConf(hadoopConf).build()
-                     )
-                   )
+      inputFile      <- path.toInputFileZIO(hadoopConf)
+      compiledFilter <- ZIO.foreach(filter) { pred =>
+                          ZIO
+                            .fromEither(Expr.compile(pred))
+                            .mapError(new IOException(_))
+                        }
+      reader         <- ZIO.fromAutoCloseable(
+                          ZIO.attemptBlockingIO {
+                            val builder = new ParquetReader.Builder(inputFile, schema, schemaEncoder)
+
+                            compiledFilter.foreach(pred => builder.withFilter(FilterCompat.get(pred)))
+                            builder.withConf(hadoopConf).build()
+                          }
+                        )
     } yield reader
 
 }
@@ -88,14 +101,17 @@ object ParquetReader {
 
   }
 
-  def configured[A <: Product: ValueDecoder](
+  def configured[A <: Product: ValueDecoder: Tag](
     hadoopConf: Configuration = new Configuration()
-  )(implicit tag: Tag[A]): ULayer[ParquetReader[A]] =
+  ): ULayer[ParquetReader[A]] =
     ZLayer.succeed(new ParquetReaderLive[A](hadoopConf))
 
-  def projected[A <: Product: ValueDecoder](
+  def projected[A <: Product: ValueDecoder: Tag](
     hadoopConf: Configuration = new Configuration()
-  )(implicit schema: Schema[A], schemaEncoder: SchemaEncoder[A], tag: Tag[A]): ULayer[ParquetReader[A]] =
+  )(implicit
+    schema: Schema[A],
+    schemaEncoder: SchemaEncoder[A]
+  ): ULayer[ParquetReader[A]] =
     ZLayer.succeed(new ParquetReaderLive[A](hadoopConf, Some(schema), Some(schemaEncoder)))
 
 }

modules/core/src/main/scala/me/mnedokushev/zio/apache/parquet/core/package.scala

@@ -1,5 +1,7 @@
 package me.mnedokushev.zio.apache.parquet
 
+import me.mnedokushev.zio.apache.parquet.core.filter.Expr
+
 package object core {
 
   val MILLIS_PER_DAY    = 86400000L
@@ -10,4 +12,8 @@ package object core {
   val DECIMAL_PRECISION = 11
   val DECIMAL_SCALE     = 2
 
+  type Lens[F, S, A]   = Expr.Column[A]
+  type Prism[F, S, A]  = Unit
+  type Traversal[S, A] = Unit
+
 }

modules/core/src/test/scala/me/mnedokushev/zio/apache/parquet/core/filter/ExprSpec.scala

@@ -0,0 +1,32 @@
+package me.mnedokushev.zio.apache.parquet.core.filter
+
+import zio.test._
+import zio._
+import zio.schema._
+import org.apache.parquet.filter2.predicate.FilterApi
+import me.mnedokushev.zio.apache.parquet.core.Value
+
+object ExprSpec extends ZIOSpecDefault {
+
+  case class MyRecord(a: String, b: Int)
+  object MyRecord {
+    implicit val recordSchema = DeriveSchema.gen[MyRecord]
+
+    val (a0, b0) = Filter.columns[MyRecord]
+  }
+
+  override def spec: Spec[TestEnvironment with Scope, Any] =
+    suite("ExprSpec")(
+      test("foo") {
+        val result   = Expr.compile(MyRecord.a0 === "bar" and MyRecord.b0 > 1)
+        val expected =
+          FilterApi.and(
+            FilterApi.eq(FilterApi.binaryColumn("a"), Value.string("bar").value),
+            FilterApi.gt(FilterApi.intColumn("b"), Int.box(Value.int(1).value))
+          )
+
+        assertTrue(result.contains(expected))
+      }
+    )
+
+}