diff --git a/.gitignore b/.gitignore
new file mode 100644
index 000000000..4eaf66513
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,10 @@
+target
+.idea
+*.iml
+derby.log
+metastore_db/
+spark-warehouse/
+dependency-reduced-pom.xml
+core/src/execution/generated
+prebuild
+.flattened-pom.xml
diff --git a/.scalafix.conf b/.scalafix.conf
new file mode 100644
index 000000000..0f23f57cd
--- /dev/null
+++ b/.scalafix.conf
@@ -0,0 +1,27 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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.
+rules = [
+ ExplicitResultTypes,
+ NoAutoTupling,
+ RemoveUnused,
+
+ DisableSyntax,
+ LeakingImplicitClassVal,
+ NoValInForComprehension,
+ ProcedureSyntax,
+ RedundantSyntax
+]
diff --git a/DEBUGGING.md b/DEBUGGING.md
new file mode 100644
index 000000000..e348b7215
--- /dev/null
+++ b/DEBUGGING.md
@@ -0,0 +1,96 @@
+
+
+# Comet Debugging Guide
+
+This HOWTO describes how to debug JVM code and Native code concurrently. The guide assumes you have:
+1. Intellij as the Java IDE
+2. CLion as the Native IDE. For Rust code, the CLion Rust language plugin is required. Note that the
+Intellij Rust plugin is not sufficient.
+3. CLion/LLDB as the native debugger. CLion ships with a bundled LLDB and the Rust community has
+its own packaging of LLDB (`lldb-rust`). Both provide a better display of Rust symbols than plain
+LLDB or the LLDB that is bundled with XCode. We will use the LLDB packaged with CLion for this guide.
+4. We will use a Comet _unit_ test as the canonical use case.
+
+_Caveat: The steps here have only been tested with JDK 11_ on Mac (M1)
+
+## Debugging for Advanced Developers
+
+Add a `.lldbinit` to comet/core. This is not strictly necessary but will be useful if you want to
+use advanced `lldb` debugging.
+
+### In Intellij
+
+1. Set a breakpoint in `NativeBase.load()`, at a point _after_ the Comet library has been loaded.
+
+1. Add a Debug Configuration for the unit test
+
+1. In the Debug Configuration for that unit test add `-Xint` as a JVM parameter. This option is
+undocumented *magic*. Without this, the LLDB debugger hits a EXC_BAD_ACCESS (or EXC_BAD_INSTRUCTION) from
+which one cannot recover.
+
+1. Add a println to the unit test to print the PID of the JVM process. (jps can also be used but this is less error prone if you have multiple jvm processes running)
+ ``` JDK8
+ println("Waiting for Debugger: PID - ", ManagementFactory.getRuntimeMXBean().getName())
+ ```
+ This will print something like : `PID@your_machine_name`.
+
+ For JDK9 and newer
+ ```JDK9
+ println("Waiting for Debugger: PID - ", ProcessHandle.current.pid)
+ ```
+
+ ==> Note the PID
+
+1. Debug-run the test in Intellij and wait for the breakpoint to be hit
+
+### In CLion
+
+1. After the breakpoint is hit in Intellij, in Clion (or LLDB from terminal or editor) -
+
+ 1. Attach to the jvm process (make sure the PID matches). In CLion, this is `Run -> Atttach to process`
+
+ 1. Put your breakpoint in the native code
+
+1. Go back to intellij and resume the process.
+
+1. Most debugging in CLion is similar to Intellij. For advanced LLDB based debugging the LLDB command line can be accessed from the LLDB tab in the Debugger view. Refer to the [LLDB manual](https://lldb.llvm.org/use/tutorial.html) for LLDB commands.
+
+### After your debugging is done,
+
+1. In CLion, detach from the process if not already detached
+
+2. In Intellij, the debugger might have lost track of the process. If so, the debugger tab
+ will show the process as running (even if the test/job is shown as completed).
+
+3. Close the debugger tab, and if the IDS asks whether it should terminate the process,
+ click Yes.
+
+4. In terminal, use jps to identify the process with the process id you were debugging. If
+ it shows up as running, kill -9 [pid]. If that doesn't remove the process, don't bother,
+ the process will be left behind as a zombie and will consume no (significant) resources.
+ Eventually it will be cleaned up when you reboot possibly after a software update.
+
+### Additional Info
+
+OpenJDK mailing list on debugging the JDK on MacOS
+https://mail.openjdk.org/pipermail/hotspot-dev/2019-September/039429.html
+
+Detecting the debugger
+https://stackoverflow.com/questions/5393403/can-a-java-application-detect-that-a-debugger-is-attached#:~:text=No.,to%20let%20your%20app%20continue.&text=I%20know%20that%20those%20are,meant%20with%20my%20first%20phrase).
diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md
new file mode 100644
index 000000000..1793bb9e2
--- /dev/null
+++ b/DEVELOPMENT.md
@@ -0,0 +1,65 @@
+
+
+# Comet Development Guide
+
+## Project Layout
+
+```
+├── common <- common Java/Scala code
+├── conf <- configuration files
+├── core <- core native code, in Rust
+├── spark <- Spark integration
+```
+
+## Development Setup
+
+1. Make sure `JAVA_HOME` is set and point to JDK 11 installation.
+2. Install Rust toolchain. The easiest way is to use
+ [rustup](https://rustup.rs).
+
+## Build & Test
+
+A few common commands are specified in project's `Makefile`:
+
+- `make`: compile the entire project, but don't run tests
+- `make test`: compile the project and run tests in both Rust and Java
+ side.
+- `make release`: compile the project and creates a release build. This
+ is useful when you want to test Comet local installation in another project
+ such as Spark.
+- `make clean`: clean up the workspace
+- `bin/comet-spark-shell -d . -o spark/target/` run Comet spark shell for V1 datasources
+- `bin/comet-spark-shell -d . -o spark/target/ --conf spark.sql.sources.useV1SourceList=""` run Comet spark shell for V2 datasources
+
+## Benchmark
+
+There's a `make` command to run micro benchmarks in the repo. For
+instance:
+
+```
+make benchmark-org.apache.spark.sql.benchmark.CometReadBenchmark
+```
+
+To run TPC-H or TPC-DS micro benchmarks, please follow the instructions
+in the respective source code, e.g., `CometTPCHQueryBenchmark`.
+
+## Debugging
+Comet is a multi-language project with native code written in Rust and JVM code written in Java and Scala.
+It is possible to debug both native and JVM code concurrently as described in the [DEBUGGING guide](DEBUGGING.md)
diff --git a/EXPRESSIONS.md b/EXPRESSIONS.md
new file mode 100644
index 000000000..40e140aba
--- /dev/null
+++ b/EXPRESSIONS.md
@@ -0,0 +1,96 @@
+
+
+# Expressions Supported by Comet
+
+The following Spark expressions are currently available:
+
++ Literals
++ Arithmetic Operators
+ + UnaryMinus
+ + Add/Minus/Multiply/Divide/Remainder
++ Conditional functions
+ + Case When
+ + If
++ Cast
++ Coalesce
++ Boolean functions
+ + And
+ + Or
+ + Not
+ + EqualTo
+ + EqualNullSafe
+ + GreaterThan
+ + GreaterThanOrEqual
+ + LessThan
+ + LessThanOrEqual
+ + IsNull
+ + IsNotNull
+ + In
++ String functions
+ + Substring
+ + Coalesce
+ + StringSpace
+ + Like
+ + Contains
+ + Startswith
+ + Endswith
+ + Ascii
+ + Bit_length
+ + Octet_length
+ + Upper
+ + Lower
+ + Chr
+ + Initcap
+ + Trim/Btrim/Ltrim/Rtrim
+ + Concat_ws
+ + Repeat
+ + Length
+ + Reverse
+ + Instr
+ + Replace
+ + Translate
++ Bitwise functions
+ + Shiftright/Shiftleft
++ Date/Time functions
+ + Year/Hour/Minute/Second
++ Math functions
+ + Abs
+ + Acos
+ + Asin
+ + Atan
+ + Atan2
+ + Cos
+ + Exp
+ + Ln
+ + Log10
+ + Log2
+ + Pow
+ + Round
+ + Signum
+ + Sin
+ + Sqrt
+ + Tan
+ + Ceil
+ + Floor
++ Aggregate functions
+ + Count
+ + Sum
+ + Max
+ + Min
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 000000000..d74c6b599
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,212 @@
+
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ 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.
+
+
+This project includes code from Apache Aurora.
+
+* dev/release/{release,changelog,release-candidate} are based on the scripts from
+ Apache Aurora
+
+Copyright: 2016 The Apache Software Foundation.
+Home page: https://aurora.apache.org/
+License: http://www.apache.org/licenses/LICENSE-2.0
diff --git a/Makefile b/Makefile
new file mode 100644
index 000000000..475c88736
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,85 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+
+.PHONY: all core jvm test clean release-linux release bench
+
+all: core jvm
+
+core:
+ cd core && cargo build
+jvm:
+ mvn clean package -DskipTests $(PROFILES)
+test:
+ mvn clean
+ # We need to compile CometException so that the cargo test can pass
+ mvn compile -pl common -DskipTests $(PROFILES)
+ cd core && cargo build && \
+ LD_LIBRARY_PATH=${LD_LIBRARY_PATH:+${LD_LIBRARY_PATH}:}${JAVA_HOME}/lib:${JAVA_HOME}/lib/server:${JAVA_HOME}/lib/jli && \
+ DYLD_LIBRARY_PATH=${DYLD_LIBRARY_PATH:+${DYLD_LIBRARY_PATH}:}${JAVA_HOME}/lib:${JAVA_HOME}/lib/server:${JAVA_HOME}/lib/jli \
+ RUST_BACKTRACE=1 cargo test
+ SPARK_HOME=`pwd` COMET_CONF_DIR=$(shell pwd)/conf RUST_BACKTRACE=1 mvn verify $(PROFILES)
+clean:
+ cd core && cargo clean
+ mvn clean
+ rm -rf .dist
+bench:
+ cd core && LD_LIBRARY_PATH=${LD_LIBRARY_PATH:+${LD_LIBRARY_PATH}:}${JAVA_HOME}/lib:${JAVA_HOME}/lib/server:${JAVA_HOME}/lib/jli && \
+ DYLD_LIBRARY_PATH=${DYLD_LIBRARY_PATH:+${DYLD_LIBRARY_PATH}:}${JAVA_HOME}/lib:${JAVA_HOME}/lib/server:${JAVA_HOME}/lib/jli \
+ RUSTFLAGS="-Ctarget-cpu=native" cargo bench $(filter-out $@,$(MAKECMDGOALS))
+format:
+ mvn compile test-compile scalafix:scalafix -Psemanticdb $(PROFILES)
+ mvn spotless:apply $(PROFILES)
+
+core-amd64:
+ rustup target add x86_64-apple-darwin
+ cd core && RUSTFLAGS="-Ctarget-cpu=skylake -Ctarget-feature=-prefer-256-bit" CC=o64-clang CXX=o64-clang++ cargo build --target x86_64-apple-darwin --features nightly --release
+ mkdir -p common/target/classes/org/apache/comet/darwin/x86_64
+ cp core/target/x86_64-apple-darwin/release/libcomet.dylib common/target/classes/org/apache/comet/darwin/x86_64
+ cd core && RUSTFLAGS="-Ctarget-cpu=haswell -Ctarget-feature=-prefer-256-bit" cargo build --features nightly --release
+ mkdir -p common/target/classes/org/apache/comet/linux/amd64
+ cp core/target/release/libcomet.so common/target/classes/org/apache/comet/linux/amd64
+ jar -cf common/target/comet-native-x86_64.jar \
+ -C common/target/classes/org/apache/comet darwin \
+ -C common/target/classes/org/apache/comet linux
+ ./dev/deploy-file common/target/comet-native-x86_64.jar comet-native-x86_64${COMET_CLASSIFIER} jar
+
+core-arm64:
+ rustup target add aarch64-apple-darwin
+ cd core && RUSTFLAGS="-Ctarget-cpu=apple-m1" CC=arm64-apple-darwin21.4-clang CXX=arm64-apple-darwin21.4-clang++ CARGO_FEATURE_NEON=1 cargo build --target aarch64-apple-darwin --features nightly --release
+ mkdir -p common/target/classes/org/apache/comet/darwin/aarch64
+ cp core/target/aarch64-apple-darwin/release/libcomet.dylib common/target/classes/org/apache/comet/darwin/aarch64
+ cd core && RUSTFLAGS="-Ctarget-cpu=native" cargo build --features nightly --release
+ mkdir -p common/target/classes/org/apache/comet/linux/aarch64
+ cp core/target/release/libcomet.so common/target/classes/org/apache/comet/linux/aarch64
+ jar -cf common/target/comet-native-aarch64.jar \
+ -C common/target/classes/org/apache/comet darwin \
+ -C common/target/classes/org/apache/comet linux
+ ./dev/deploy-file common/target/comet-native-aarch64.jar comet-native-aarch64${COMET_CLASSIFIER} jar
+
+release-linux: clean
+ rustup target add aarch64-apple-darwin x86_64-apple-darwin
+ cd core && RUSTFLAGS="-Ctarget-cpu=apple-m1" CC=arm64-apple-darwin21.4-clang CXX=arm64-apple-darwin21.4-clang++ CARGO_FEATURE_NEON=1 cargo build --target aarch64-apple-darwin --features nightly --release
+ cd core && RUSTFLAGS="-Ctarget-cpu=skylake -Ctarget-feature=-prefer-256-bit" CC=o64-clang CXX=o64-clang++ cargo build --target x86_64-apple-darwin --features nightly --release
+ cd core && RUSTFLAGS="-Ctarget-cpu=native -Ctarget-feature=-prefer-256-bit" cargo build --features nightly --release
+ mvn install -Prelease -DskipTests $(PROFILES)
+release:
+ cd core && RUSTFLAGS="-Ctarget-cpu=native" cargo build --features nightly --release
+ mvn install -Prelease -DskipTests $(PROFILES)
+benchmark-%: clean release
+ cd spark && COMET_CONF_DIR=$(shell pwd)/conf MAVEN_OPTS='-Xmx20g' .mvn exec:java -Dexec.mainClass="$*" -Dexec.classpathScope="test" -Dexec.cleanupDaemonThreads="false" -Dexec.args="$(filter-out $@,$(MAKECMDGOALS))" $(PROFILES)
+.DEFAULT:
+ @: # ignore arguments provided to benchmarks e.g. "make benchmark-foo -- --bar", we do not want to treat "--bar" as target
diff --git a/README.md b/README.md
index ba486a23f..ba50f673a 100644
--- a/README.md
+++ b/README.md
@@ -1 +1,60 @@
-# arrow-datafusion-comet
+
+
+# Apache Arrow DataFusion Comet
+
+Comet is an Apache Spark plugin that uses [Apache Arrow DataFusion](https://arrow.apache.org/datafusion/)
+as native runtime to achieve improvement in terms of query efficiency and query runtime.
+
+On a high level, Comet aims to support:
+- a native Parquet implementation, including both reader and writer
+- full implementation of Spark operators, including
+ Filter/Project/Aggregation/Join/Exchange etc.
+- full implementation of Spark built-in expressions
+- a UDF framework for users to migrate their existing UDF to native
+
+The following diagram illustrates the architecture of Comet:
+
+
+
+## Current Status
+
+The project is currently integrated into Apache Spark 3.2, 3.3, and 3.4.
+
+## Feature Parity with Apache Spark
+
+The project strives to keep feature parity with Apache Spark, that is,
+users should expect the same behavior (w.r.t features, configurations,
+query results, etc) with Comet turned on or turned off in their Spark
+jobs. In addition, Comet extension should automatically detect unsupported
+features and fallback to Spark engine.
+
+To achieve this, besides unit tests within Comet itself, we also re-use
+Spark SQL tests and make sure they all pass with Comet extension
+enabled.
+
+## Supported Platforms
+
+Linux, Apple OSX (Intel and M1)
+
+## Requirements
+
+- Apache Spark 3.2, 3.3, or 3.4
+- JDK 8 and up
+- GLIBC 2.17 (Centos 7) and up
diff --git a/bin/comet-spark-shell b/bin/comet-spark-shell
new file mode 100755
index 000000000..9ae55a1b3
--- /dev/null
+++ b/bin/comet-spark-shell
@@ -0,0 +1,84 @@
+#!/bin/bash
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+
+
+set -e
+
+if [[ -z ${SPARK_HOME} ]]; then
+ echo "ERROR: SPARK_HOME is not set: please set it to point to a Spark 3.4.0 distribution"
+ exit 1
+fi
+
+POSITIONAL=()
+while [[ $# -gt 0 ]]; do
+ key="$1"
+
+ case $key in
+ -r|--rebuild)
+ REBUILD=YES
+ shift
+ ;;
+ -o|--outdir)
+ COMET_OUTDIR=$2
+ shift
+ shift
+ ;;
+ -d|--comet-dir)
+ COMET_DIR=$2
+ shift
+ shift
+ ;;
+ *)
+ POSITIONAL+=("$1")
+ shift
+ ;;
+ esac
+done
+
+set -- "${POSITIONAL[@]}"
+
+COMET_DIR="${COMET_DIR:-$HOME/git/comet}"
+SPARK_VERSION_SHORT=3.4
+SCALA_BINARY_VERSION=2.12
+COMET_VERSION=0.1.0-SNAPSHOT
+COMET_SPARK_JAR=comet-spark-spark${SPARK_VERSION_SHORT}_${SCALA_BINARY_VERSION}-${COMET_VERSION}.jar
+COMET_OUTDIR="${COMET_OUTDIR:-/tmp}"
+
+if [[ ! -d $COMET_DIR ]]; then
+ echo "Comet repo: $COMET_DIR, doesn't exist"
+ exit 1
+fi
+
+if [[ ! -d $COMET_OUTDIR ]]; then
+ echo "Output directory for Comet Spark library: $COMET_OUTDIR, doesn't exist"
+ exit 1
+fi
+
+if [[ "X$REBUILD" == "XYES" ]]; then
+ cd $COMET_DIR && make release
+ cd $COMET_DIR/spark && cp target/${COMET_SPARK_JAR} $COMET_OUTDIR/${COMET_SPARK_JAR}
+fi
+
+RUST_BACKTRACE=1 $SPARK_HOME/bin/spark-shell \
+ --jars $COMET_OUTDIR/${COMET_SPARK_JAR} \
+ --conf spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions \
+ --conf spark.comet.enabled=true \
+ --conf spark.comet.exec.enabled=true \
+ --conf spark.comet.exec.all.enabled=true \
+$@
diff --git a/common/pom.xml b/common/pom.xml
new file mode 100644
index 000000000..f885c346d
--- /dev/null
+++ b/common/pom.xml
@@ -0,0 +1,200 @@
+
+
+
+
+
+
+ 4.0.0
+
+ org.apache.comet
+ comet-parent-spark${spark.version.short}_${scala.binary.version}
+ 0.1.0-SNAPSHOT
+ ../pom.xml
+
+
+ comet-common-spark${spark.version.short}_${scala.binary.version}
+ comet-common
+
+
+
+ false
+
+
+
+
+ org.apache.spark
+ spark-sql_${scala.binary.version}
+
+
+ org.apache.parquet
+ parquet-column
+
+
+ org.apache.parquet
+ parquet-hadoop
+
+
+ org.apache.arrow
+ arrow-vector
+
+
+ org.apache.arrow
+ arrow-memory-unsafe
+
+
+ org.apache.arrow
+ arrow-c-data
+
+
+ junit
+ junit
+ test
+
+
+ org.assertj
+ assertj-core
+ test
+
+
+
+
+
+
+ io.github.git-commit-id
+ git-commit-id-maven-plugin
+ 5.0.0
+
+
+ get-the-git-infos
+
+ revision
+
+ initialize
+
+
+
+ true
+ ${project.build.outputDirectory}/comet-git-info.properties
+ full
+
+ ^git.branch$
+ ^git.build.*$
+ ^git.commit.id.(abbrev|full)$
+ ^git.remote.*$
+
+
+
+
+ org.apache.maven.plugins
+ maven-shade-plugin
+
+
+ package
+
+ shade
+
+
+ true
+ true
+ false
+ true
+
+
+
+ org.apache.arrow:*
+
+
+
+
+ *:*
+
+ **/*.thrift
+ git.properties
+ log4j.properties
+ log4j2.properties
+ arrow-git.properties
+
+
+
+ org.apache.arrow:arrow-vector
+
+
+ codegen/**
+
+
+
+
+
+ org.apache.arrow
+ ${comet.shade.packageName}.arrow
+
+
+ org/apache/arrow/c/jni/JniWrapper
+ org/apache/arrow/c/jni/PrivateData
+ org/apache/arrow/c/jni/CDataJniException
+
+ org/apache/arrow/c/ArrayStreamExporter$ExportedArrayStreamPrivateData
+
+
+
+
+
+
+
+
+
+ net.alchim31.maven
+ scala-maven-plugin
+
+
+
+
+ ${project.basedir}/src/main/resources
+
+
+ ${project.basedir}/../core/target/x86_64-apple-darwin/release
+
+ libcomet.dylib
+
+ org/apache/comet/darwin/x86_64
+
+
+ ${project.basedir}/../core/target/aarch64-apple-darwin/release
+
+ libcomet.dylib
+
+ org/apache/comet/darwin/aarch64
+
+
+ ${jni.dir}
+
+ libcomet.dylib
+ libcomet.so
+
+ org/apache/comet/${platform}/${arch}
+
+
+
+
+
diff --git a/common/src/main/java/org/apache/arrow/c/ArrowImporter.java b/common/src/main/java/org/apache/arrow/c/ArrowImporter.java
new file mode 100644
index 000000000..90398cb72
--- /dev/null
+++ b/common/src/main/java/org/apache/arrow/c/ArrowImporter.java
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.arrow.c;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.types.pojo.Field;
+
+/**
+ * This class is used to import Arrow schema and array from native execution/shuffle. We cannot use
+ * Arrow's Java API to import schema and array directly because Arrow's Java API `Data.importField`
+ * initiates a new `SchemaImporter` for each field. Each `SchemaImporter` maintains an internal
+ * dictionary id counter. So the dictionary ids for multiple dictionary columns will conflict with
+ * each other and cause data corruption.
+ */
+public class ArrowImporter {
+ private final SchemaImporter importer;
+ private final BufferAllocator allocator;
+
+ public ArrowImporter(BufferAllocator allocator) {
+ this.allocator = allocator;
+ this.importer = new SchemaImporter(allocator);
+ }
+
+ Field importField(ArrowSchema schema, CDataDictionaryProvider provider) {
+ Field var4;
+ try {
+ var4 = importer.importField(schema, provider);
+ } finally {
+ schema.release();
+ schema.close();
+ }
+
+ return var4;
+ }
+
+ public FieldVector importVector(
+ ArrowArray array, ArrowSchema schema, CDataDictionaryProvider provider) {
+ Field field = importField(schema, provider);
+ FieldVector vector = field.createVector(allocator);
+ Data.importIntoVector(allocator, array, vector, provider);
+ return vector;
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/CometRuntimeException.java b/common/src/main/java/org/apache/comet/CometRuntimeException.java
new file mode 100644
index 000000000..b136ad165
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/CometRuntimeException.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet;
+
+/** The parent class for all Comet runtime exceptions */
+public class CometRuntimeException extends RuntimeException {
+ public CometRuntimeException(String message) {
+ super(message);
+ }
+
+ public CometRuntimeException(String message, Throwable cause) {
+ super(message, cause);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/NativeBase.java b/common/src/main/java/org/apache/comet/NativeBase.java
new file mode 100644
index 000000000..42357b9bf
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/NativeBase.java
@@ -0,0 +1,278 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FilenameFilter;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.file.Files;
+import java.nio.file.StandardCopyOption;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.spark.sql.comet.util.Utils;
+
+import static org.apache.comet.Constants.LOG_CONF_NAME;
+import static org.apache.comet.Constants.LOG_CONF_PATH;
+
+/** Base class for JNI bindings. MUST be inherited by all classes that introduce JNI APIs. */
+public abstract class NativeBase {
+ static final String ARROW_UNSAFE_MEMORY_ACCESS = "arrow.enable_unsafe_memory_access";
+ static final String ARROW_NULL_CHECK_FOR_GET = "arrow.enable_null_check_for_get";
+
+ private static final Logger LOG = LoggerFactory.getLogger(NativeBase.class);
+ private static final String NATIVE_LIB_NAME = "comet";
+
+ private static final String libraryToLoad = System.mapLibraryName(NATIVE_LIB_NAME);
+ private static boolean loaded = false;
+ private static final String searchPattern = "libcomet-";
+
+ static {
+ if (!isLoaded()) {
+ load();
+ }
+ }
+
+ public static synchronized boolean isLoaded() {
+ return loaded;
+ }
+
+ // Only for testing
+ static synchronized void setLoaded(boolean b) {
+ loaded = b;
+ }
+
+ static synchronized void load() {
+ if (loaded) {
+ return;
+ }
+
+ cleanupOldTempLibs();
+
+ // Check if the arch used by JDK is the same as arch on the host machine, in particular,
+ // whether x86_64 JDK is used in arm64 Mac
+ if (!checkArch()) {
+ LOG.warn(
+ "Comet is disabled. JDK compiled for x86_64 is used in a Mac based on Apple Silicon. "
+ + "In order to use Comet, Please install a JDK version for ARM64 architecture");
+ return;
+ }
+
+ // Try to load Comet library from the java.library.path.
+ try {
+ System.loadLibrary(libraryToLoad);
+ loaded = true;
+ } catch (UnsatisfiedLinkError ex) {
+ // Doesn't exist, so proceed to loading bundled library.
+ bundleLoadLibrary();
+ }
+
+ initWithLogConf();
+ // Only set the Arrow properties when debugging mode is off
+ if (!(boolean) CometConf.COMET_DEBUG_ENABLED().get()) {
+ setArrowProperties();
+ }
+ }
+
+ /**
+ * Use the bundled native libraries. Functionally equivalent to System.loadLibrary.
+ */
+ private static void bundleLoadLibrary() {
+ String resourceName = resourceName();
+ InputStream is = NativeBase.class.getResourceAsStream(resourceName);
+ if (is == null) {
+ throw new UnsupportedOperationException(
+ "Unsupported OS/arch, cannot find "
+ + resourceName
+ + ". Please try building from source.");
+ }
+
+ File tempLib = null;
+ File tempLibLock = null;
+ try {
+ // Create the .lck file first to avoid a race condition
+ // with other concurrently running Java processes using Comet.
+ tempLibLock = File.createTempFile(searchPattern, "." + os().libExtension + ".lck");
+ tempLib = new File(tempLibLock.getAbsolutePath().replaceFirst(".lck$", ""));
+ // copy to tempLib
+ Files.copy(is, tempLib.toPath(), StandardCopyOption.REPLACE_EXISTING);
+ System.load(tempLib.getAbsolutePath());
+ loaded = true;
+ } catch (IOException e) {
+ throw new IllegalStateException("Cannot unpack libcomet: " + e);
+ } finally {
+ if (!loaded) {
+ if (tempLib != null && tempLib.exists()) {
+ if (!tempLib.delete()) {
+ LOG.error(
+ "Cannot unpack libcomet / cannot delete a temporary native library " + tempLib);
+ }
+ }
+ if (tempLibLock != null && tempLibLock.exists()) {
+ if (!tempLibLock.delete()) {
+ LOG.error(
+ "Cannot unpack libcomet / cannot delete a temporary lock file " + tempLibLock);
+ }
+ }
+ } else {
+ tempLib.deleteOnExit();
+ tempLibLock.deleteOnExit();
+ }
+ }
+ }
+
+ private static void initWithLogConf() {
+ String logConfPath = System.getProperty(LOG_CONF_PATH(), Utils.getConfPath(LOG_CONF_NAME()));
+
+ // If both the system property and the environmental variable failed to find a log
+ // configuration, then fall back to using the deployed default
+ if (logConfPath == null) {
+ LOG.info(
+ "Couldn't locate log file from either COMET_CONF_DIR or comet.log.file.path. "
+ + "Using default log configuration which emits to stdout");
+ logConfPath = "";
+ } else {
+ LOG.info("Using {} for native library logging", logConfPath);
+ }
+ init(logConfPath);
+ }
+
+ private static void cleanupOldTempLibs() {
+ String tempFolder = new File(System.getProperty("java.io.tmpdir")).getAbsolutePath();
+ File dir = new File(tempFolder);
+
+ File[] tempLibFiles =
+ dir.listFiles(
+ new FilenameFilter() {
+ public boolean accept(File dir, String name) {
+ return name.startsWith(searchPattern) && !name.endsWith(".lck");
+ }
+ });
+
+ if (tempLibFiles != null) {
+ for (File tempLibFile : tempLibFiles) {
+ File lckFile = new File(tempLibFile.getAbsolutePath() + ".lck");
+ if (!lckFile.exists()) {
+ try {
+ tempLibFile.delete();
+ } catch (SecurityException e) {
+ LOG.error("Failed to delete old temp lib", e);
+ }
+ }
+ }
+ }
+ }
+
+ // Set Arrow related properties upon initializing native, such as enabling unsafe memory access
+ // as well as disabling null check for get, for performance reasons.
+ private static void setArrowProperties() {
+ setPropertyIfNull(ARROW_UNSAFE_MEMORY_ACCESS, "true");
+ setPropertyIfNull(ARROW_NULL_CHECK_FOR_GET, "false");
+ }
+
+ private static void setPropertyIfNull(String key, String value) {
+ if (System.getProperty(key) == null) {
+ LOG.info("Setting system property {} to {}", key, value);
+ System.setProperty(key, value);
+ } else {
+ LOG.info(
+ "Skip setting system property {} to {}, because it is already set to {}",
+ key,
+ value,
+ System.getProperty(key));
+ }
+ }
+
+ private enum OS {
+ // Even on Windows, the default compiler from cpptasks (gcc) uses .so as a shared lib extension
+ WINDOWS("win32", "so"),
+ LINUX("linux", "so"),
+ MAC("darwin", "dylib"),
+ SOLARIS("solaris", "so");
+ public final String name, libExtension;
+
+ OS(String name, String libExtension) {
+ this.name = name;
+ this.libExtension = libExtension;
+ }
+ }
+
+ private static String arch() {
+ return System.getProperty("os.arch");
+ }
+
+ private static OS os() {
+ String osName = System.getProperty("os.name");
+ if (osName.contains("Linux")) {
+ return OS.LINUX;
+ } else if (osName.contains("Mac")) {
+ return OS.MAC;
+ } else if (osName.contains("Windows")) {
+ return OS.WINDOWS;
+ } else if (osName.contains("Solaris") || osName.contains("SunOS")) {
+ return OS.SOLARIS;
+ } else {
+ throw new UnsupportedOperationException("Unsupported operating system: " + osName);
+ }
+ }
+
+ // For some reason users will get JVM crash when running Comet that is compiled for `aarch64`
+ // using a JVM that is compiled against `amd64`. Here we check if that is the case and fallback
+ // to Spark accordingly.
+ private static boolean checkArch() {
+ if (os() == OS.MAC) {
+ try {
+ String javaArch = arch();
+ Process process = Runtime.getRuntime().exec("uname -a");
+ if (process.waitFor() == 0) {
+ BufferedReader in = new BufferedReader(new InputStreamReader(process.getInputStream()));
+ String line;
+ while ((line = in.readLine()) != null) {
+ if (javaArch.equals("x86_64") && line.contains("ARM64")) {
+ return false;
+ }
+ }
+ }
+ } catch (IOException | InterruptedException e) {
+ LOG.warn("Error parsing host architecture", e);
+ }
+ }
+
+ return true;
+ }
+
+ private static String resourceName() {
+ OS os = os();
+ String packagePrefix = NativeBase.class.getPackage().getName().replace('.', '/');
+
+ return "/" + packagePrefix + "/" + os.name + "/" + arch() + "/" + libraryToLoad;
+ }
+
+ /**
+ * Initialize the native library through JNI.
+ *
+ * @param logConfPath location to the native log configuration file
+ */
+ static native void init(String logConfPath);
+}
diff --git a/common/src/main/java/org/apache/comet/ParquetRuntimeException.java b/common/src/main/java/org/apache/comet/ParquetRuntimeException.java
new file mode 100644
index 000000000..1f81587e7
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/ParquetRuntimeException.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet;
+
+/** The parent class for the subset of Comet runtime exceptions related to Parquet. */
+public class ParquetRuntimeException extends CometRuntimeException {
+ public ParquetRuntimeException(String message) {
+ super(message);
+ }
+
+ public ParquetRuntimeException(String message, Throwable cause) {
+ super(message, cause);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/AbstractColumnReader.java b/common/src/main/java/org/apache/comet/parquet/AbstractColumnReader.java
new file mode 100644
index 000000000..099c7b973
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/AbstractColumnReader.java
@@ -0,0 +1,116 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.TimestampNTZType$;
+
+import org.apache.comet.CometConf;
+import org.apache.comet.vector.CometVector;
+
+/** Base class for Comet Parquet column reader implementations. */
+public abstract class AbstractColumnReader implements AutoCloseable {
+ protected static final Logger LOG = LoggerFactory.getLogger(AbstractColumnReader.class);
+
+ /** The Spark data type. */
+ protected final DataType type;
+
+ /** Parquet column descriptor. */
+ protected final ColumnDescriptor descriptor;
+
+ /**
+ * Whether to always return 128 bit decimals, regardless of its precision. If false, this will
+ * return 32, 64 or 128 bit decimals depending on the precision.
+ */
+ protected final boolean useDecimal128;
+
+ /**
+ * Whether to return dates/timestamps that were written with legacy hybrid (Julian + Gregorian)
+ * calendar as it is. If this is true, Comet will return them as it is, instead of rebasing them
+ * to the new Proleptic Gregorian calendar. If this is false, Comet will throw exceptions when
+ * seeing these dates/timestamps.
+ */
+ protected final boolean useLegacyDateTimestamp;
+
+ /** The size of one batch, gets updated by 'readBatch' */
+ protected int batchSize;
+
+ /** A pointer to the native implementation of ColumnReader. */
+ protected long nativeHandle;
+
+ public AbstractColumnReader(
+ DataType type,
+ ColumnDescriptor descriptor,
+ boolean useDecimal128,
+ boolean useLegacyDateTimestamp) {
+ this.type = type;
+ this.descriptor = descriptor;
+ this.useDecimal128 = useDecimal128;
+ this.useLegacyDateTimestamp = useLegacyDateTimestamp;
+ TypeUtil.checkParquetType(descriptor, type);
+ }
+
+ public ColumnDescriptor getDescriptor() {
+ return descriptor;
+ }
+
+ /**
+ * Set the batch size of this reader to be 'batchSize'. Also initializes the native column reader.
+ */
+ public void setBatchSize(int batchSize) {
+ assert nativeHandle == 0
+ : "Native column reader shouldn't be initialized before " + "'setBatchSize' is called";
+ this.batchSize = batchSize;
+ initNative();
+ }
+
+ /**
+ * Reads a batch of 'total' new rows.
+ *
+ * @param total the total number of rows to read
+ */
+ public abstract void readBatch(int total);
+
+ /** Returns the {@link CometVector} read by this reader. */
+ public abstract CometVector currentBatch();
+
+ @Override
+ public void close() {
+ if (nativeHandle != 0) {
+ LOG.debug("Closing the column reader");
+ Native.closeColumnReader(nativeHandle);
+ nativeHandle = 0;
+ }
+ }
+
+ protected void initNative() {
+ LOG.debug("initializing the native column reader");
+ DataType readType = (boolean) CometConf.COMET_SCHEMA_EVOLUTION_ENABLED().get() ? type : null;
+ boolean useLegacyDateTimestampOrNTZ =
+ useLegacyDateTimestamp || type == TimestampNTZType$.MODULE$;
+ nativeHandle =
+ Utils.initColumnReader(
+ descriptor, readType, batchSize, useDecimal128, useLegacyDateTimestampOrNTZ);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/BatchReader.java b/common/src/main/java/org/apache/comet/parquet/BatchReader.java
new file mode 100644
index 000000000..87302b372
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/BatchReader.java
@@ -0,0 +1,620 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import scala.Option;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.mapreduce.InputSplit;
+import org.apache.hadoop.mapreduce.RecordReader;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.apache.parquet.HadoopReadOptions;
+import org.apache.parquet.ParquetReadOptions;
+import org.apache.parquet.Preconditions;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.page.PageReadStore;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.apache.parquet.schema.MessageType;
+import org.apache.parquet.schema.Type;
+import org.apache.spark.TaskContext;
+import org.apache.spark.TaskContext$;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.comet.parquet.CometParquetReadSupport;
+import org.apache.spark.sql.execution.datasources.PartitionedFile;
+import org.apache.spark.sql.execution.datasources.parquet.ParquetToSparkSchemaConverter;
+import org.apache.spark.sql.execution.metric.SQLMetric;
+import org.apache.spark.sql.types.*;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+import org.apache.spark.util.AccumulatorV2;
+
+import org.apache.comet.CometConf;
+import org.apache.comet.shims.ShimBatchReader;
+import org.apache.comet.shims.ShimFileFormat;
+import org.apache.comet.vector.CometVector;
+
+/**
+ * A vectorized Parquet reader that reads a Parquet file in a batched fashion.
+ *
+ *
Example of how to use this:
+ *
+ *
+ * BatchReader reader = new BatchReader(parquetFile, batchSize);
+ * try {
+ * reader.init();
+ * while (reader.readBatch()) {
+ * ColumnarBatch batch = reader.currentBatch();
+ * // consume the batch
+ * }
+ * } finally { // resources associated with the reader should be released
+ * reader.close();
+ * }
+ *
+ */
+public class BatchReader extends RecordReader implements Closeable {
+ private static final Logger LOG = LoggerFactory.getLogger(FileReader.class);
+
+ private Configuration conf;
+ private int capacity;
+ private boolean isCaseSensitive;
+ private boolean useFieldId;
+ private boolean ignoreMissingIds;
+ private StructType partitionSchema;
+ private InternalRow partitionValues;
+ private PartitionedFile file;
+ private final Map metrics;
+
+ private long rowsRead;
+ private StructType sparkSchema;
+ private MessageType requestedSchema;
+ private CometVector[] vectors;
+ private AbstractColumnReader[] columnReaders;
+ private ColumnarBatch currentBatch;
+ private Future> prefetchTask;
+ private LinkedBlockingQueue> prefetchQueue;
+ private FileReader fileReader;
+ private boolean[] missingColumns;
+ private boolean isInitialized;
+ private ParquetMetadata footer;
+
+ /** The total number of rows across all row groups of the input split. */
+ private long totalRowCount;
+
+ /**
+ * The total number of rows loaded so far, including all the rows from row groups that we've
+ * processed and the current row group.
+ */
+ private long totalRowsLoaded;
+
+ /**
+ * Whether the native scan should always return decimal represented by 128 bits, regardless of its
+ * precision. Normally, this should be true if native execution is enabled, since Arrow compute
+ * kernels doesn't support 32 and 64 bit decimals yet.
+ */
+ private boolean useDecimal128;
+
+ /** Whether to use the lazy materialization reader for reading columns. */
+ private boolean useLazyMaterialization;
+
+ /**
+ * Whether to return dates/timestamps that were written with legacy hybrid (Julian + Gregorian)
+ * calendar as it is. If this is true, Comet will return them as it is, instead of rebasing them
+ * to the new Proleptic Gregorian calendar. If this is false, Comet will throw exceptions when
+ * seeing these dates/timestamps.
+ */
+ private boolean useLegacyDateTimestamp;
+
+ /** The TaskContext object for executing this task. */
+ private final TaskContext taskContext;
+
+ // Only for testing
+ public BatchReader(String file, int capacity) {
+ this(file, capacity, null, null);
+ }
+
+ // Only for testing
+ public BatchReader(
+ String file, int capacity, StructType partitionSchema, InternalRow partitionValues) {
+ this(new Configuration(), file, capacity, partitionSchema, partitionValues);
+ }
+
+ // Only for testing
+ public BatchReader(
+ Configuration conf,
+ String file,
+ int capacity,
+ StructType partitionSchema,
+ InternalRow partitionValues) {
+ conf.set("spark.sql.parquet.binaryAsString", "false");
+ conf.set("spark.sql.parquet.int96AsTimestamp", "false");
+ conf.set("spark.sql.caseSensitive", "false");
+ conf.set("spark.sql.parquet.inferTimestampNTZ.enabled", "true");
+ conf.set("spark.sql.legacy.parquet.nanosAsLong", "false");
+
+ this.conf = conf;
+ this.capacity = capacity;
+ this.isCaseSensitive = false;
+ this.useFieldId = false;
+ this.ignoreMissingIds = false;
+ this.partitionSchema = partitionSchema;
+ this.partitionValues = partitionValues;
+
+ this.file = ShimBatchReader.newPartitionedFile(partitionValues, file);
+ this.metrics = new HashMap<>();
+
+ this.taskContext = TaskContext$.MODULE$.get();
+ }
+
+ public BatchReader(AbstractColumnReader[] columnReaders) {
+ // Todo: set useDecimal128 and useLazyMaterialization
+ int numColumns = columnReaders.length;
+ this.columnReaders = new AbstractColumnReader[numColumns];
+ vectors = new CometVector[numColumns];
+ currentBatch = new ColumnarBatch(vectors);
+ // This constructor is used by Iceberg only. The columnReaders are
+ // initialized in Iceberg, so no need to call the init()
+ isInitialized = true;
+ this.taskContext = TaskContext$.MODULE$.get();
+ this.metrics = new HashMap<>();
+ }
+
+ BatchReader(
+ Configuration conf,
+ PartitionedFile inputSplit,
+ ParquetMetadata footer,
+ int capacity,
+ StructType sparkSchema,
+ boolean isCaseSensitive,
+ boolean useFieldId,
+ boolean ignoreMissingIds,
+ boolean useLegacyDateTimestamp,
+ StructType partitionSchema,
+ InternalRow partitionValues,
+ Map metrics) {
+ this.conf = conf;
+ this.capacity = capacity;
+ this.sparkSchema = sparkSchema;
+ this.isCaseSensitive = isCaseSensitive;
+ this.useFieldId = useFieldId;
+ this.ignoreMissingIds = ignoreMissingIds;
+ this.useLegacyDateTimestamp = useLegacyDateTimestamp;
+ this.partitionSchema = partitionSchema;
+ this.partitionValues = partitionValues;
+ this.file = inputSplit;
+ this.footer = footer;
+ this.metrics = metrics;
+ this.taskContext = TaskContext$.MODULE$.get();
+ }
+
+ /**
+ * Initialize this reader. The reason we don't do it in the constructor is that we want to close
+ * any resource hold by this reader when error happens during the initialization.
+ */
+ public void init() throws URISyntaxException, IOException {
+ useDecimal128 =
+ conf.getBoolean(
+ CometConf.COMET_USE_DECIMAL_128().key(),
+ (Boolean) CometConf.COMET_USE_DECIMAL_128().defaultValue().get());
+ useLazyMaterialization =
+ conf.getBoolean(
+ CometConf.COMET_USE_LAZY_MATERIALIZATION().key(),
+ (Boolean) CometConf.COMET_USE_LAZY_MATERIALIZATION().defaultValue().get());
+
+ long start = file.start();
+ long length = file.length();
+ String filePath = file.filePath().toString();
+
+ ParquetReadOptions.Builder builder = HadoopReadOptions.builder(conf, new Path(filePath));
+
+ if (start >= 0 && length >= 0) {
+ builder = builder.withRange(start, start + length);
+ }
+ ParquetReadOptions readOptions = builder.build();
+
+ // TODO: enable off-heap buffer when they are ready
+ ReadOptions cometReadOptions = ReadOptions.builder(conf).build();
+
+ Path path = new Path(new URI(filePath));
+ fileReader =
+ new FileReader(
+ CometInputFile.fromPath(path, conf), footer, readOptions, cometReadOptions, metrics);
+ requestedSchema = fileReader.getFileMetaData().getSchema();
+ MessageType fileSchema = requestedSchema;
+
+ if (sparkSchema == null) {
+ sparkSchema = new ParquetToSparkSchemaConverter(conf).convert(requestedSchema);
+ } else {
+ requestedSchema =
+ CometParquetReadSupport.clipParquetSchema(
+ requestedSchema, sparkSchema, isCaseSensitive, useFieldId, ignoreMissingIds);
+ if (requestedSchema.getColumns().size() != sparkSchema.size()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Spark schema has %d columns while " + "Parquet schema has %d columns",
+ sparkSchema.size(), requestedSchema.getColumns().size()));
+ }
+ }
+
+ totalRowCount = fileReader.getRecordCount();
+ List columns = requestedSchema.getColumns();
+ int numColumns = columns.size();
+ if (partitionSchema != null) numColumns += partitionSchema.size();
+ columnReaders = new AbstractColumnReader[numColumns];
+
+ // Initialize missing columns and use null vectors for them
+ missingColumns = new boolean[columns.size()];
+ List paths = requestedSchema.getPaths();
+ StructField[] nonPartitionFields = sparkSchema.fields();
+ for (int i = 0; i < requestedSchema.getFieldCount(); i++) {
+ Type t = requestedSchema.getFields().get(i);
+ Preconditions.checkState(
+ t.isPrimitive() && !t.isRepetition(Type.Repetition.REPEATED),
+ "Complex type is not supported");
+ String[] colPath = paths.get(i);
+ if (nonPartitionFields[i].name().equals(ShimFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME())) {
+ // Values of ROW_INDEX_TEMPORARY_COLUMN_NAME column are always populated with
+ // generated row indexes, rather than read from the file.
+ // TODO(SPARK-40059): Allow users to include columns named
+ // FileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME in their schemas.
+ long[] rowIndices = fileReader.getRowIndices();
+ columnReaders[i] = new RowIndexColumnReader(nonPartitionFields[i], capacity, rowIndices);
+ missingColumns[i] = true;
+ } else if (fileSchema.containsPath(colPath)) {
+ ColumnDescriptor fd = fileSchema.getColumnDescription(colPath);
+ if (!fd.equals(columns.get(i))) {
+ throw new UnsupportedOperationException("Schema evolution is not supported");
+ }
+ missingColumns[i] = false;
+ } else {
+ if (columns.get(i).getMaxDefinitionLevel() == 0) {
+ throw new IOException(
+ "Required column '"
+ + Arrays.toString(colPath)
+ + "' is missing"
+ + " in data file "
+ + filePath);
+ }
+ ConstantColumnReader reader =
+ new ConstantColumnReader(nonPartitionFields[i], capacity, useDecimal128);
+ columnReaders[i] = reader;
+ missingColumns[i] = true;
+ }
+ }
+
+ // Initialize constant readers for partition columns
+ if (partitionSchema != null) {
+ StructField[] partitionFields = partitionSchema.fields();
+ for (int i = columns.size(); i < columnReaders.length; i++) {
+ int fieldIndex = i - columns.size();
+ StructField field = partitionFields[fieldIndex];
+ ConstantColumnReader reader =
+ new ConstantColumnReader(field, capacity, partitionValues, fieldIndex, useDecimal128);
+ columnReaders[i] = reader;
+ }
+ }
+
+ vectors = new CometVector[numColumns];
+ currentBatch = new ColumnarBatch(vectors);
+ fileReader.setRequestedSchema(requestedSchema.getColumns());
+
+ // For test purpose only
+ // If the last external accumulator is `NumRowGroupsAccumulator`, the row group number to read
+ // will be updated to the accumulator. So we can check if the row groups are filtered or not
+ // in test case.
+ // Note that this tries to get thread local TaskContext object, if this is called at other
+ // thread, it won't update the accumulator.
+ if (taskContext != null) {
+ Option> accu = taskContext.taskMetrics().externalAccums().lastOption();
+ if (accu.isDefined() && accu.get().getClass().getSimpleName().equals("NumRowGroupsAcc")) {
+ @SuppressWarnings("unchecked")
+ AccumulatorV2 intAccum = (AccumulatorV2) accu.get();
+ intAccum.add(fileReader.getRowGroups().size());
+ }
+ }
+
+ // Pre-fetching
+ boolean preFetchEnabled =
+ conf.getBoolean(
+ CometConf.COMET_SCAN_PREFETCH_ENABLED().key(),
+ (boolean) CometConf.COMET_SCAN_PREFETCH_ENABLED().defaultValue().get());
+
+ if (preFetchEnabled) {
+ LOG.info("Prefetch enabled for BatchReader.");
+ this.prefetchQueue = new LinkedBlockingQueue<>();
+ }
+
+ isInitialized = true;
+ synchronized (this) {
+ // if prefetch is enabled, `init()` is called in separate thread. When
+ // `BatchReader.nextBatch()` is called asynchronously, it is possibly that
+ // `init()` is not called or finished. We need to hold on `nextBatch` until
+ // initialization of `BatchReader` is done. Once we are close to finish
+ // initialization, we notify the waiting thread of `nextBatch` to continue.
+ notifyAll();
+ }
+ }
+
+ public void setSparkSchema(StructType schema) {
+ this.sparkSchema = schema;
+ }
+
+ public AbstractColumnReader[] getColumnReaders() {
+ return columnReaders;
+ }
+
+ @Override
+ public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext)
+ throws IOException, InterruptedException {
+ // Do nothing. The initialization work is done in 'init' already.
+ }
+
+ @Override
+ public boolean nextKeyValue() throws IOException {
+ return nextBatch();
+ }
+
+ @Override
+ public Void getCurrentKey() {
+ return null;
+ }
+
+ @Override
+ public ColumnarBatch getCurrentValue() {
+ return currentBatch();
+ }
+
+ @Override
+ public float getProgress() {
+ return (float) rowsRead / totalRowCount;
+ }
+
+ /**
+ * Returns the current columnar batch being read.
+ *
+ *
Note that this must be called AFTER {@link BatchReader#nextBatch()}.
+ */
+ public ColumnarBatch currentBatch() {
+ return currentBatch;
+ }
+
+ // Only for testing
+ public Future> getPrefetchTask() {
+ return this.prefetchTask;
+ }
+
+ // Only for testing
+ public LinkedBlockingQueue> getPrefetchQueue() {
+ return this.prefetchQueue;
+ }
+
+ /**
+ * Loads the next batch of rows.
+ *
+ * @return true if there are no more rows to read, false otherwise.
+ */
+ public boolean nextBatch() throws IOException {
+ if (this.prefetchTask == null) {
+ Preconditions.checkState(isInitialized, "init() should be called first!");
+ } else {
+ // If prefetch is enabled, this reader will be initialized asynchronously from a
+ // different thread. Wait until it is initialized
+ while (!isInitialized) {
+ synchronized (this) {
+ try {
+ // Wait until initialization of current `BatchReader` is finished (i.e., `init()`),
+ // is done. It is possibly that `init()` is done after entering this while loop,
+ // so a short timeout is given.
+ wait(100);
+
+ // Checks if prefetch task is finished. If so, tries to get exception if any.
+ if (prefetchTask.isDone()) {
+ Option exception = prefetchTask.get();
+ if (exception.isDefined()) {
+ throw exception.get();
+ }
+ }
+ } catch (RuntimeException e) {
+ // Spark will check certain exception e.g. `SchemaColumnConvertNotSupportedException`.
+ throw e;
+ } catch (Throwable e) {
+ throw new IOException(e);
+ }
+ }
+ }
+ }
+
+ if (rowsRead >= totalRowCount) return false;
+ boolean hasMore;
+
+ try {
+ hasMore = loadNextRowGroupIfNecessary();
+ } catch (RuntimeException e) {
+ // Spark will check certain exception e.g. `SchemaColumnConvertNotSupportedException`.
+ throw e;
+ } catch (Throwable e) {
+ throw new IOException(e);
+ }
+
+ if (!hasMore) return false;
+ int batchSize = (int) Math.min(capacity, totalRowsLoaded - rowsRead);
+
+ return nextBatch(batchSize);
+ }
+
+ public boolean nextBatch(int batchSize) {
+ long totalDecodeTime = 0, totalLoadTime = 0;
+ for (int i = 0; i < columnReaders.length; i++) {
+ AbstractColumnReader reader = columnReaders[i];
+ long startNs = System.nanoTime();
+ reader.readBatch(batchSize);
+ totalDecodeTime += System.nanoTime() - startNs;
+ startNs = System.nanoTime();
+ vectors[i] = reader.currentBatch();
+ totalLoadTime += System.nanoTime() - startNs;
+ }
+
+ SQLMetric decodeMetric = metrics.get("ParquetNativeDecodeTime");
+ if (decodeMetric != null) {
+ decodeMetric.add(totalDecodeTime);
+ }
+ SQLMetric loadMetric = metrics.get("ParquetNativeLoadTime");
+ if (loadMetric != null) {
+ loadMetric.add(totalLoadTime);
+ }
+
+ currentBatch.setNumRows(batchSize);
+ rowsRead += batchSize;
+ return true;
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (columnReaders != null) {
+ for (AbstractColumnReader reader : columnReaders) {
+ if (reader != null) {
+ reader.close();
+ }
+ }
+ }
+ if (fileReader != null) {
+ fileReader.close();
+ fileReader = null;
+ }
+ }
+
+ private boolean loadNextRowGroupIfNecessary() throws Throwable {
+ // More rows can be read from loaded row group. No need to load next one.
+ if (rowsRead != totalRowsLoaded) return true;
+
+ SQLMetric rowGroupTimeMetric = metrics.get("ParquetLoadRowGroupTime");
+ SQLMetric numRowGroupsMetric = metrics.get("ParquetRowGroups");
+ long startNs = System.nanoTime();
+
+ PageReadStore rowGroupReader = null;
+ if (prefetchTask != null && prefetchQueue != null) {
+ // Wait for pre-fetch task to finish.
+ Pair rowGroupReaderPair = prefetchQueue.take();
+ rowGroupReader = rowGroupReaderPair.getLeft();
+
+ // Update incremental byte read metric. Because this metric in Spark is maintained
+ // by thread local variable, we need to manually update it.
+ // TODO: We may expose metrics from `FileReader` and get from it directly.
+ long incBytesRead = rowGroupReaderPair.getRight();
+ FileSystem.getAllStatistics().stream()
+ .forEach(statistic -> statistic.incrementBytesRead(incBytesRead));
+ } else {
+ rowGroupReader = fileReader.readNextRowGroup();
+ }
+
+ if (rowGroupTimeMetric != null) {
+ rowGroupTimeMetric.add(System.nanoTime() - startNs);
+ }
+ if (rowGroupReader == null) {
+ return false;
+ }
+ if (numRowGroupsMetric != null) {
+ numRowGroupsMetric.add(1);
+ }
+
+ List columns = requestedSchema.getColumns();
+ for (int i = 0; i < columns.size(); i++) {
+ if (missingColumns[i]) continue;
+ if (columnReaders[i] != null) columnReaders[i].close();
+ // TODO: handle tz, datetime & int96 rebase
+ // TODO: consider passing page reader via ctor - however we need to fix the shading issue
+ // from Iceberg side.
+ DataType dataType = sparkSchema.fields()[i].dataType();
+ ColumnReader reader =
+ Utils.getColumnReader(
+ dataType,
+ columns.get(i),
+ capacity,
+ useDecimal128,
+ useLazyMaterialization,
+ useLegacyDateTimestamp);
+ reader.setPageReader(rowGroupReader.getPageReader(columns.get(i)));
+ columnReaders[i] = reader;
+ }
+ totalRowsLoaded += rowGroupReader.getRowCount();
+ return true;
+ }
+
+ // Submits a prefetch task for this reader.
+ public void submitPrefetchTask(ExecutorService threadPool) {
+ this.prefetchTask = threadPool.submit(new PrefetchTask());
+ }
+
+ // A task for prefetching parquet row groups.
+ private class PrefetchTask implements Callable> {
+ private long getBytesRead() {
+ return FileSystem.getAllStatistics().stream()
+ .mapToLong(s -> s.getThreadStatistics().getBytesRead())
+ .sum();
+ }
+
+ @Override
+ public Option call() throws Exception {
+ // Gets the bytes read so far.
+ long baseline = getBytesRead();
+
+ try {
+ init();
+
+ while (true) {
+ PageReadStore rowGroupReader = fileReader.readNextRowGroup();
+
+ if (rowGroupReader == null) {
+ // Reaches the end of row groups.
+ return Option.empty();
+ } else {
+ long incBytesRead = getBytesRead() - baseline;
+
+ prefetchQueue.add(Pair.of(rowGroupReader, incBytesRead));
+ }
+ }
+ } catch (Throwable e) {
+ // Returns exception thrown from the reader. The reader will re-throw it.
+ return Option.apply(e);
+ } finally {
+ if (fileReader != null) {
+ fileReader.closeStream();
+ }
+ }
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/BloomFilterReader.java b/common/src/main/java/org/apache/comet/parquet/BloomFilterReader.java
new file mode 100644
index 000000000..a23216c7f
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/BloomFilterReader.java
@@ -0,0 +1,253 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.parquet.column.values.bloomfilter.BlockSplitBloomFilter;
+import org.apache.parquet.column.values.bloomfilter.BloomFilter;
+import org.apache.parquet.crypto.AesCipher;
+import org.apache.parquet.crypto.InternalColumnDecryptionSetup;
+import org.apache.parquet.crypto.InternalFileDecryptor;
+import org.apache.parquet.crypto.ModuleCipherFactory;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+import org.apache.parquet.filter2.predicate.FilterPredicate;
+import org.apache.parquet.filter2.predicate.Operators;
+import org.apache.parquet.filter2.predicate.UserDefinedPredicate;
+import org.apache.parquet.format.BlockCipher;
+import org.apache.parquet.format.BloomFilterHeader;
+import org.apache.parquet.format.Util;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.io.SeekableInputStream;
+
+public class BloomFilterReader implements FilterPredicate.Visitor {
+ private static final Logger LOG = LoggerFactory.getLogger(BloomFilterReader.class);
+ private static final boolean BLOCK_MIGHT_MATCH = false;
+ private static final boolean BLOCK_CANNOT_MATCH = true;
+
+ private final Map columns;
+ private final Map cache = new HashMap<>();
+ private final InternalFileDecryptor fileDecryptor;
+ private final SeekableInputStream inputStream;
+
+ BloomFilterReader(
+ BlockMetaData block, InternalFileDecryptor fileDecryptor, SeekableInputStream inputStream) {
+ this.columns = new HashMap<>();
+ for (ColumnChunkMetaData column : block.getColumns()) {
+ columns.put(column.getPath(), column);
+ }
+ this.fileDecryptor = fileDecryptor;
+ this.inputStream = inputStream;
+ }
+
+ @Override
+ public > Boolean visit(Operators.Eq eq) {
+ T value = eq.getValue();
+
+ if (value == null) {
+ // the bloom filter bitset contains only non-null values so isn't helpful. this
+ // could check the column stats, but the StatisticsFilter is responsible
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ Operators.Column filterColumn = eq.getColumn();
+ ColumnChunkMetaData meta = columns.get(filterColumn.getColumnPath());
+ if (meta == null) {
+ // the column isn't in this file so all values are null, but the value
+ // must be non-null because of the above check.
+ return BLOCK_CANNOT_MATCH;
+ }
+
+ try {
+ BloomFilter bloomFilter = readBloomFilter(meta);
+ if (bloomFilter != null && !bloomFilter.findHash(bloomFilter.hash(value))) {
+ return BLOCK_CANNOT_MATCH;
+ }
+ } catch (RuntimeException e) {
+ LOG.warn(e.getMessage());
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ @Override
+ public > Boolean visit(Operators.NotEq notEq) {
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ @Override
+ public > Boolean visit(Operators.Lt lt) {
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ @Override
+ public > Boolean visit(Operators.LtEq ltEq) {
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ @Override
+ public > Boolean visit(Operators.Gt gt) {
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ @Override
+ public > Boolean visit(Operators.GtEq gtEq) {
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ @Override
+ public Boolean visit(Operators.And and) {
+ return and.getLeft().accept(this) || and.getRight().accept(this);
+ }
+
+ @Override
+ public Boolean visit(Operators.Or or) {
+ return or.getLeft().accept(this) && or.getRight().accept(this);
+ }
+
+ @Override
+ public Boolean visit(Operators.Not not) {
+ throw new IllegalArgumentException(
+ "This predicate "
+ + not
+ + " contains a not! Did you forget"
+ + " to run this predicate through LogicalInverseRewriter?");
+ }
+
+ @Override
+ public , U extends UserDefinedPredicate> Boolean visit(
+ Operators.UserDefined udp) {
+ return visit(udp, false);
+ }
+
+ @Override
+ public , U extends UserDefinedPredicate> Boolean visit(
+ Operators.LogicalNotUserDefined udp) {
+ return visit(udp.getUserDefined(), true);
+ }
+
+ private , U extends UserDefinedPredicate> Boolean visit(
+ Operators.UserDefined ud, boolean inverted) {
+ return BLOCK_MIGHT_MATCH;
+ }
+
+ BloomFilter readBloomFilter(ColumnChunkMetaData meta) {
+ if (cache.containsKey(meta.getPath())) {
+ return cache.get(meta.getPath());
+ }
+ try {
+ if (!cache.containsKey(meta.getPath())) {
+ BloomFilter bloomFilter = readBloomFilterInternal(meta);
+ if (bloomFilter == null) {
+ return null;
+ }
+
+ cache.put(meta.getPath(), bloomFilter);
+ }
+ return cache.get(meta.getPath());
+ } catch (IOException e) {
+ LOG.error("Failed to read Bloom filter data", e);
+ }
+
+ return null;
+ }
+
+ private BloomFilter readBloomFilterInternal(ColumnChunkMetaData meta) throws IOException {
+ long bloomFilterOffset = meta.getBloomFilterOffset();
+ if (bloomFilterOffset < 0) {
+ return null;
+ }
+
+ // Prepare to decrypt Bloom filter (for encrypted columns)
+ BlockCipher.Decryptor bloomFilterDecryptor = null;
+ byte[] bloomFilterHeaderAAD = null;
+ byte[] bloomFilterBitsetAAD = null;
+ if (null != fileDecryptor && !fileDecryptor.plaintextFile()) {
+ InternalColumnDecryptionSetup columnDecryptionSetup =
+ fileDecryptor.getColumnSetup(meta.getPath());
+ if (columnDecryptionSetup.isEncrypted()) {
+ bloomFilterDecryptor = columnDecryptionSetup.getMetaDataDecryptor();
+ bloomFilterHeaderAAD =
+ AesCipher.createModuleAAD(
+ fileDecryptor.getFileAAD(),
+ ModuleCipherFactory.ModuleType.BloomFilterHeader,
+ meta.getRowGroupOrdinal(),
+ columnDecryptionSetup.getOrdinal(),
+ -1);
+ bloomFilterBitsetAAD =
+ AesCipher.createModuleAAD(
+ fileDecryptor.getFileAAD(),
+ ModuleCipherFactory.ModuleType.BloomFilterBitset,
+ meta.getRowGroupOrdinal(),
+ columnDecryptionSetup.getOrdinal(),
+ -1);
+ }
+ }
+
+ // Read Bloom filter data header.
+ inputStream.seek(bloomFilterOffset);
+ BloomFilterHeader bloomFilterHeader;
+ try {
+ bloomFilterHeader =
+ Util.readBloomFilterHeader(inputStream, bloomFilterDecryptor, bloomFilterHeaderAAD);
+ } catch (IOException e) {
+ LOG.warn("read no bloom filter");
+ return null;
+ }
+
+ int numBytes = bloomFilterHeader.getNumBytes();
+ if (numBytes <= 0 || numBytes > BlockSplitBloomFilter.UPPER_BOUND_BYTES) {
+ LOG.warn("the read bloom filter size is wrong, size is {}", bloomFilterHeader.getNumBytes());
+ return null;
+ }
+
+ if (!bloomFilterHeader.getHash().isSetXXHASH()
+ || !bloomFilterHeader.getAlgorithm().isSetBLOCK()
+ || !bloomFilterHeader.getCompression().isSetUNCOMPRESSED()) {
+ LOG.warn(
+ "the read bloom filter is not supported yet, algorithm = {}, hash = {}, "
+ + "compression = {}",
+ bloomFilterHeader.getAlgorithm(),
+ bloomFilterHeader.getHash(),
+ bloomFilterHeader.getCompression());
+ return null;
+ }
+
+ byte[] bitset;
+ if (null == bloomFilterDecryptor) {
+ bitset = new byte[numBytes];
+ inputStream.readFully(bitset);
+ } else {
+ bitset = bloomFilterDecryptor.decrypt(inputStream, bloomFilterBitsetAAD);
+ if (bitset.length != numBytes) {
+ throw new ParquetCryptoRuntimeException("Wrong length of decrypted bloom filter bitset");
+ }
+ }
+ return new BlockSplitBloomFilter(bitset);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/ColumnIndexReader.java b/common/src/main/java/org/apache/comet/parquet/ColumnIndexReader.java
new file mode 100644
index 000000000..4e4f6ba0d
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/ColumnIndexReader.java
@@ -0,0 +1,230 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.parquet.crypto.AesCipher;
+import org.apache.parquet.crypto.InternalColumnDecryptionSetup;
+import org.apache.parquet.crypto.InternalFileDecryptor;
+import org.apache.parquet.crypto.ModuleCipherFactory;
+import org.apache.parquet.format.BlockCipher;
+import org.apache.parquet.format.Util;
+import org.apache.parquet.format.converter.ParquetMetadataConverter;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.internal.column.columnindex.ColumnIndex;
+import org.apache.parquet.internal.column.columnindex.OffsetIndex;
+import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore;
+import org.apache.parquet.internal.hadoop.metadata.IndexReference;
+import org.apache.parquet.io.SeekableInputStream;
+
+class ColumnIndexReader implements ColumnIndexStore {
+ private static final Logger LOG = LoggerFactory.getLogger(ColumnIndexReader.class);
+
+ // Used for columns are not in this parquet file
+ private static final IndexStore MISSING_INDEX_STORE =
+ new IndexStore() {
+ @Override
+ public ColumnIndex getColumnIndex() {
+ return null;
+ }
+
+ @Override
+ public OffsetIndex getOffsetIndex() {
+ return null;
+ }
+ };
+
+ private static final ColumnIndexReader EMPTY =
+ new ColumnIndexReader(new BlockMetaData(), Collections.emptySet(), null, null) {
+ @Override
+ public ColumnIndex getColumnIndex(ColumnPath column) {
+ return null;
+ }
+
+ @Override
+ public OffsetIndex getOffsetIndex(ColumnPath column) {
+ throw new MissingOffsetIndexException(column);
+ }
+ };
+
+ private final InternalFileDecryptor fileDecryptor;
+ private final SeekableInputStream inputStream;
+ private final Map store;
+
+ /**
+ * Creates a column index store which lazily reads column/offset indexes for the columns in paths.
+ * Paths are the set of columns used for the projection.
+ */
+ static ColumnIndexReader create(
+ BlockMetaData block,
+ Set paths,
+ InternalFileDecryptor fileDecryptor,
+ SeekableInputStream inputStream) {
+ try {
+ return new ColumnIndexReader(block, paths, fileDecryptor, inputStream);
+ } catch (MissingOffsetIndexException e) {
+ return EMPTY;
+ }
+ }
+
+ private ColumnIndexReader(
+ BlockMetaData block,
+ Set paths,
+ InternalFileDecryptor fileDecryptor,
+ SeekableInputStream inputStream) {
+ this.fileDecryptor = fileDecryptor;
+ this.inputStream = inputStream;
+ Map store = new HashMap<>();
+ for (ColumnChunkMetaData column : block.getColumns()) {
+ ColumnPath path = column.getPath();
+ if (paths.contains(path)) {
+ store.put(path, new IndexStoreImpl(column));
+ }
+ }
+ this.store = store;
+ }
+
+ @Override
+ public ColumnIndex getColumnIndex(ColumnPath column) {
+ return store.getOrDefault(column, MISSING_INDEX_STORE).getColumnIndex();
+ }
+
+ @Override
+ public OffsetIndex getOffsetIndex(ColumnPath column) {
+ return store.getOrDefault(column, MISSING_INDEX_STORE).getOffsetIndex();
+ }
+
+ private interface IndexStore {
+ ColumnIndex getColumnIndex();
+
+ OffsetIndex getOffsetIndex();
+ }
+
+ private class IndexStoreImpl implements IndexStore {
+ private final ColumnChunkMetaData meta;
+ private ColumnIndex columnIndex;
+ private boolean columnIndexRead;
+ private final OffsetIndex offsetIndex;
+
+ IndexStoreImpl(ColumnChunkMetaData meta) {
+ this.meta = meta;
+ OffsetIndex oi;
+ try {
+ oi = readOffsetIndex(meta);
+ } catch (IOException e) {
+ // If the I/O issue still stands it will fail the reading later;
+ // otherwise we fail the filtering only with a missing offset index.
+ LOG.warn("Unable to read offset index for column {}", meta.getPath(), e);
+ oi = null;
+ }
+ if (oi == null) {
+ throw new MissingOffsetIndexException(meta.getPath());
+ }
+ offsetIndex = oi;
+ }
+
+ @Override
+ public ColumnIndex getColumnIndex() {
+ if (!columnIndexRead) {
+ try {
+ columnIndex = readColumnIndex(meta);
+ } catch (IOException e) {
+ // If the I/O issue still stands it will fail the reading later;
+ // otherwise we fail the filtering only with a missing column index.
+ LOG.warn("Unable to read column index for column {}", meta.getPath(), e);
+ }
+ columnIndexRead = true;
+ }
+ return columnIndex;
+ }
+
+ @Override
+ public OffsetIndex getOffsetIndex() {
+ return offsetIndex;
+ }
+ }
+
+ // Visible for testing
+ ColumnIndex readColumnIndex(ColumnChunkMetaData column) throws IOException {
+ IndexReference ref = column.getColumnIndexReference();
+ if (ref == null) {
+ return null;
+ }
+ inputStream.seek(ref.getOffset());
+
+ BlockCipher.Decryptor columnIndexDecryptor = null;
+ byte[] columnIndexAAD = null;
+ if (null != fileDecryptor && !fileDecryptor.plaintextFile()) {
+ InternalColumnDecryptionSetup columnDecryptionSetup =
+ fileDecryptor.getColumnSetup(column.getPath());
+ if (columnDecryptionSetup.isEncrypted()) {
+ columnIndexDecryptor = columnDecryptionSetup.getMetaDataDecryptor();
+ columnIndexAAD =
+ AesCipher.createModuleAAD(
+ fileDecryptor.getFileAAD(),
+ ModuleCipherFactory.ModuleType.ColumnIndex,
+ column.getRowGroupOrdinal(),
+ columnDecryptionSetup.getOrdinal(),
+ -1);
+ }
+ }
+ return ParquetMetadataConverter.fromParquetColumnIndex(
+ column.getPrimitiveType(),
+ Util.readColumnIndex(inputStream, columnIndexDecryptor, columnIndexAAD));
+ }
+
+ // Visible for testing
+ OffsetIndex readOffsetIndex(ColumnChunkMetaData column) throws IOException {
+ IndexReference ref = column.getOffsetIndexReference();
+ if (ref == null) {
+ return null;
+ }
+ inputStream.seek(ref.getOffset());
+
+ BlockCipher.Decryptor offsetIndexDecryptor = null;
+ byte[] offsetIndexAAD = null;
+ if (null != fileDecryptor && !fileDecryptor.plaintextFile()) {
+ InternalColumnDecryptionSetup columnDecryptionSetup =
+ fileDecryptor.getColumnSetup(column.getPath());
+ if (columnDecryptionSetup.isEncrypted()) {
+ offsetIndexDecryptor = columnDecryptionSetup.getMetaDataDecryptor();
+ offsetIndexAAD =
+ AesCipher.createModuleAAD(
+ fileDecryptor.getFileAAD(),
+ ModuleCipherFactory.ModuleType.OffsetIndex,
+ column.getRowGroupOrdinal(),
+ columnDecryptionSetup.getOrdinal(),
+ -1);
+ }
+ }
+ return ParquetMetadataConverter.fromParquetOffsetIndex(
+ Util.readOffsetIndex(inputStream, offsetIndexDecryptor, offsetIndexAAD));
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/ColumnPageReader.java b/common/src/main/java/org/apache/comet/parquet/ColumnPageReader.java
new file mode 100644
index 000000000..744d12830
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/ColumnPageReader.java
@@ -0,0 +1,252 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.List;
+import java.util.Queue;
+
+import org.apache.parquet.bytes.BytesInput;
+import org.apache.parquet.column.page.DataPage;
+import org.apache.parquet.column.page.DataPageV1;
+import org.apache.parquet.column.page.DataPageV2;
+import org.apache.parquet.column.page.DictionaryPage;
+import org.apache.parquet.column.page.PageReader;
+import org.apache.parquet.compression.CompressionCodecFactory;
+import org.apache.parquet.crypto.AesCipher;
+import org.apache.parquet.crypto.ModuleCipherFactory;
+import org.apache.parquet.format.BlockCipher;
+import org.apache.parquet.internal.column.columnindex.OffsetIndex;
+import org.apache.parquet.io.ParquetDecodingException;
+
+public class ColumnPageReader implements PageReader {
+ private final CompressionCodecFactory.BytesInputDecompressor decompressor;
+ private final long valueCount;
+ private final Queue compressedPages;
+ private final DictionaryPage compressedDictionaryPage;
+
+ private final OffsetIndex offsetIndex;
+ private final long rowCount;
+ private int pageIndex = 0;
+
+ private final BlockCipher.Decryptor blockDecryptor;
+ private final byte[] dataPageAAD;
+ private final byte[] dictionaryPageAAD;
+
+ ColumnPageReader(
+ CompressionCodecFactory.BytesInputDecompressor decompressor,
+ List compressedPages,
+ DictionaryPage compressedDictionaryPage,
+ OffsetIndex offsetIndex,
+ long rowCount,
+ BlockCipher.Decryptor blockDecryptor,
+ byte[] fileAAD,
+ int rowGroupOrdinal,
+ int columnOrdinal) {
+ this.decompressor = decompressor;
+ this.compressedPages = new ArrayDeque<>(compressedPages);
+ this.compressedDictionaryPage = compressedDictionaryPage;
+ long count = 0;
+ for (DataPage p : compressedPages) {
+ count += p.getValueCount();
+ }
+ this.valueCount = count;
+ this.offsetIndex = offsetIndex;
+ this.rowCount = rowCount;
+ this.blockDecryptor = blockDecryptor;
+
+ if (blockDecryptor != null) {
+ dataPageAAD =
+ AesCipher.createModuleAAD(
+ fileAAD, ModuleCipherFactory.ModuleType.DataPage, rowGroupOrdinal, columnOrdinal, 0);
+ dictionaryPageAAD =
+ AesCipher.createModuleAAD(
+ fileAAD,
+ ModuleCipherFactory.ModuleType.DictionaryPage,
+ rowGroupOrdinal,
+ columnOrdinal,
+ -1);
+ } else {
+ dataPageAAD = null;
+ dictionaryPageAAD = null;
+ }
+ }
+
+ @Override
+ public long getTotalValueCount() {
+ return valueCount;
+ }
+
+ /** Returns the total value count of the current page. */
+ public int getPageValueCount() {
+ return compressedPages.element().getValueCount();
+ }
+
+ /** Skips the current page so it won't be returned by {@link #readPage()} */
+ public void skipPage() {
+ compressedPages.poll();
+ pageIndex++;
+ }
+
+ @Override
+ public DataPage readPage() {
+ final DataPage compressedPage = compressedPages.poll();
+ if (compressedPage == null) {
+ return null;
+ }
+ final int currentPageIndex = pageIndex++;
+
+ if (null != blockDecryptor) {
+ AesCipher.quickUpdatePageAAD(dataPageAAD, getPageOrdinal(currentPageIndex));
+ }
+
+ return compressedPage.accept(
+ new DataPage.Visitor() {
+ @Override
+ public DataPage visit(DataPageV1 dataPageV1) {
+ try {
+ BytesInput bytes = dataPageV1.getBytes();
+ if (null != blockDecryptor) {
+ bytes = BytesInput.from(blockDecryptor.decrypt(bytes.toByteArray(), dataPageAAD));
+ }
+ BytesInput decompressed =
+ decompressor.decompress(bytes, dataPageV1.getUncompressedSize());
+
+ final DataPageV1 decompressedPage;
+ if (offsetIndex == null) {
+ decompressedPage =
+ new DataPageV1(
+ decompressed,
+ dataPageV1.getValueCount(),
+ dataPageV1.getUncompressedSize(),
+ dataPageV1.getStatistics(),
+ dataPageV1.getRlEncoding(),
+ dataPageV1.getDlEncoding(),
+ dataPageV1.getValueEncoding());
+ } else {
+ long firstRowIndex = offsetIndex.getFirstRowIndex(currentPageIndex);
+ decompressedPage =
+ new DataPageV1(
+ decompressed,
+ dataPageV1.getValueCount(),
+ dataPageV1.getUncompressedSize(),
+ firstRowIndex,
+ Math.toIntExact(
+ offsetIndex.getLastRowIndex(currentPageIndex, rowCount)
+ - firstRowIndex
+ + 1),
+ dataPageV1.getStatistics(),
+ dataPageV1.getRlEncoding(),
+ dataPageV1.getDlEncoding(),
+ dataPageV1.getValueEncoding());
+ }
+ if (dataPageV1.getCrc().isPresent()) {
+ decompressedPage.setCrc(dataPageV1.getCrc().getAsInt());
+ }
+ return decompressedPage;
+ } catch (IOException e) {
+ throw new ParquetDecodingException("could not decompress page", e);
+ }
+ }
+
+ @Override
+ public DataPage visit(DataPageV2 dataPageV2) {
+ if (!dataPageV2.isCompressed() && offsetIndex == null && null == blockDecryptor) {
+ return dataPageV2;
+ }
+ BytesInput pageBytes = dataPageV2.getData();
+
+ if (null != blockDecryptor) {
+ try {
+ pageBytes =
+ BytesInput.from(blockDecryptor.decrypt(pageBytes.toByteArray(), dataPageAAD));
+ } catch (IOException e) {
+ throw new ParquetDecodingException(
+ "could not convert page ByteInput to byte array", e);
+ }
+ }
+ if (dataPageV2.isCompressed()) {
+ int uncompressedSize =
+ Math.toIntExact(
+ dataPageV2.getUncompressedSize()
+ - dataPageV2.getDefinitionLevels().size()
+ - dataPageV2.getRepetitionLevels().size());
+ try {
+ pageBytes = decompressor.decompress(pageBytes, uncompressedSize);
+ } catch (IOException e) {
+ throw new ParquetDecodingException("could not decompress page", e);
+ }
+ }
+
+ if (offsetIndex == null) {
+ return DataPageV2.uncompressed(
+ dataPageV2.getRowCount(),
+ dataPageV2.getNullCount(),
+ dataPageV2.getValueCount(),
+ dataPageV2.getRepetitionLevels(),
+ dataPageV2.getDefinitionLevels(),
+ dataPageV2.getDataEncoding(),
+ pageBytes,
+ dataPageV2.getStatistics());
+ } else {
+ return DataPageV2.uncompressed(
+ dataPageV2.getRowCount(),
+ dataPageV2.getNullCount(),
+ dataPageV2.getValueCount(),
+ offsetIndex.getFirstRowIndex(currentPageIndex),
+ dataPageV2.getRepetitionLevels(),
+ dataPageV2.getDefinitionLevels(),
+ dataPageV2.getDataEncoding(),
+ pageBytes,
+ dataPageV2.getStatistics());
+ }
+ }
+ });
+ }
+
+ @Override
+ public DictionaryPage readDictionaryPage() {
+ if (compressedDictionaryPage == null) {
+ return null;
+ }
+ try {
+ BytesInput bytes = compressedDictionaryPage.getBytes();
+ if (null != blockDecryptor) {
+ bytes = BytesInput.from(blockDecryptor.decrypt(bytes.toByteArray(), dictionaryPageAAD));
+ }
+ DictionaryPage decompressedPage =
+ new DictionaryPage(
+ decompressor.decompress(bytes, compressedDictionaryPage.getUncompressedSize()),
+ compressedDictionaryPage.getDictionarySize(),
+ compressedDictionaryPage.getEncoding());
+ if (compressedDictionaryPage.getCrc().isPresent()) {
+ decompressedPage.setCrc(compressedDictionaryPage.getCrc().getAsInt());
+ }
+ return decompressedPage;
+ } catch (IOException e) {
+ throw new ParquetDecodingException("Could not decompress dictionary page", e);
+ }
+ }
+
+ private int getPageOrdinal(int currentPageIndex) {
+ return offsetIndex == null ? currentPageIndex : offsetIndex.getPageOrdinal(currentPageIndex);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/ColumnReader.java b/common/src/main/java/org/apache/comet/parquet/ColumnReader.java
new file mode 100644
index 000000000..7e45f4f9a
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/ColumnReader.java
@@ -0,0 +1,314 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.arrow.c.ArrowArray;
+import org.apache.arrow.c.ArrowSchema;
+import org.apache.arrow.c.CDataDictionaryProvider;
+import org.apache.arrow.c.Data;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.Encoding;
+import org.apache.parquet.column.page.DataPage;
+import org.apache.parquet.column.page.DataPageV1;
+import org.apache.parquet.column.page.DataPageV2;
+import org.apache.parquet.column.page.DictionaryPage;
+import org.apache.parquet.column.page.PageReader;
+import org.apache.spark.sql.types.DataType;
+
+import org.apache.comet.CometConf;
+import org.apache.comet.vector.CometDecodedVector;
+import org.apache.comet.vector.CometDictionary;
+import org.apache.comet.vector.CometDictionaryVector;
+import org.apache.comet.vector.CometPlainVector;
+import org.apache.comet.vector.CometVector;
+
+public class ColumnReader extends AbstractColumnReader {
+ protected static final Logger LOG = LoggerFactory.getLogger(ColumnReader.class);
+ protected static final BufferAllocator ALLOCATOR = new RootAllocator();
+
+ /**
+ * The current Comet vector holding all the values read by this column reader. Owned by this
+ * reader and MUST be closed after use.
+ */
+ private CometDecodedVector currentVector;
+
+ /** Dictionary values for this column. Only set if the column is using dictionary encoding. */
+ protected CometDictionary dictionary;
+
+ /** Reader for dictionary & data pages in the current column chunk. */
+ protected PageReader pageReader;
+
+ /** Whether the first data page has been loaded. */
+ private boolean firstPageLoaded = false;
+
+ /**
+ * The number of nulls in the current batch, used when we are skipping importing of Arrow vectors,
+ * in which case we'll simply update the null count of the existing vectors.
+ */
+ int currentNumNulls;
+
+ /**
+ * The number of values in the current batch, used when we are skipping importing of Arrow
+ * vectors, in which case we'll simply update the null count of the existing vectors.
+ */
+ int currentNumValues;
+
+ /**
+ * Whether the last loaded vector contains any null value. This is used to determine if we can
+ * skip vector reloading. If the flag is false, Arrow C API will skip to import the validity
+ * buffer, and therefore we cannot skip vector reloading.
+ */
+ boolean hadNull;
+
+ /** Dictionary provider for this column. */
+ private final CDataDictionaryProvider dictionaryProvider = new CDataDictionaryProvider();
+
+ public ColumnReader(
+ DataType type,
+ ColumnDescriptor descriptor,
+ int batchSize,
+ boolean useDecimal128,
+ boolean useLegacyDateTimestamp) {
+ super(type, descriptor, useDecimal128, useLegacyDateTimestamp);
+ assert batchSize > 0 : "Batch size must be positive, found " + batchSize;
+ this.batchSize = batchSize;
+ initNative();
+ }
+
+ /**
+ * Set the page reader for a new column chunk to read. Expects to call `readBatch` after this.
+ *
+ * @param pageReader the page reader for the new column chunk
+ */
+ public void setPageReader(PageReader pageReader) throws IOException {
+ this.pageReader = pageReader;
+
+ DictionaryPage dictionaryPage = pageReader.readDictionaryPage();
+ if (dictionaryPage != null) {
+ LOG.debug("dictionary page encoding = {}", dictionaryPage.getEncoding());
+ Native.setDictionaryPage(
+ nativeHandle,
+ dictionaryPage.getDictionarySize(),
+ dictionaryPage.getBytes().toByteArray(),
+ dictionaryPage.getEncoding().ordinal());
+ }
+ }
+
+ @Override
+ public void readBatch(int total) {
+ LOG.debug("Start to batch of size = " + total);
+
+ if (!firstPageLoaded) {
+ readPage();
+ firstPageLoaded = true;
+ }
+
+ // Now first reset the current columnar batch so that it can be used to fill in a new batch
+ // of values. Then, keep reading more data pages (via 'readBatch') until the current batch is
+ // full, or we have read 'total' number of values.
+ Native.resetBatch(nativeHandle);
+
+ int left = total, nullsRead = 0;
+ while (left > 0) {
+ int[] array = Native.readBatch(nativeHandle, left);
+ int valuesRead = array[0];
+ nullsRead += array[1];
+ if (valuesRead < left) {
+ readPage();
+ }
+ left -= valuesRead;
+ }
+
+ this.currentNumValues = total;
+ this.currentNumNulls = nullsRead;
+ }
+
+ /** Returns the {@link CometVector} read by this reader. */
+ @Override
+ public CometVector currentBatch() {
+ return loadVector();
+ }
+
+ @Override
+ public void close() {
+ if (currentVector != null) {
+ currentVector.close();
+ currentVector = null;
+ }
+ dictionaryProvider.close();
+ super.close();
+ }
+
+ /** Returns a decoded {@link CometDecodedVector Comet vector}. */
+ public CometDecodedVector loadVector() {
+ // Only re-use Comet vector iff:
+ // 1. if we're not using dictionary encoding, since with dictionary encoding, the native
+ // side may fallback to plain encoding and the underlying memory address for the vector
+ // will change as result.
+ // 2. if the column type is of fixed width, in other words, string/binary are not supported
+ // since the native side may resize the vector and therefore change memory address.
+ // 3. if the last loaded vector contains null values: if values of last vector are all not
+ // null, Arrow C data API will skip loading the native validity buffer, therefore we
+ // should not re-use the vector in that case.
+ // 4. if the last loaded vector doesn't contain any null value, but the current vector also
+ // are all not null, which means we can also re-use the loaded vector.
+ // 5. if the new number of value is the same or smaller
+ if ((hadNull || currentNumNulls == 0)
+ && currentVector != null
+ && dictionary == null
+ && currentVector.isFixedLength()
+ && currentVector.numValues() >= currentNumValues) {
+ currentVector.setNumNulls(currentNumNulls);
+ currentVector.setNumValues(currentNumValues);
+ return currentVector;
+ }
+
+ LOG.debug("Reloading vector");
+
+ // Close the previous vector first to release struct memory allocated to import Arrow array &
+ // schema from native side, through the C data interface
+ if (currentVector != null) {
+ currentVector.close();
+ }
+
+ long[] addresses = Native.currentBatch(nativeHandle);
+
+ try (ArrowArray array = ArrowArray.wrap(addresses[0]);
+ ArrowSchema schema = ArrowSchema.wrap(addresses[1])) {
+ FieldVector vector = Data.importVector(ALLOCATOR, array, schema, dictionaryProvider);
+ DictionaryEncoding dictionaryEncoding = vector.getField().getDictionary();
+
+ CometPlainVector cometVector = new CometPlainVector(vector, useDecimal128);
+
+ // Update whether the current vector contains any null values. This is used in the following
+ // batch(s) to determine whether we can skip loading the native vector.
+ hadNull = cometVector.hasNull();
+
+ if (dictionaryEncoding == null) {
+ if (dictionary != null) {
+ // This means the column was using dictionary encoding but now has fall-back to plain
+ // encoding, on the native side. Setting 'dictionary' to null here, so we can use it as
+ // a condition to check if we can re-use vector later.
+ dictionary = null;
+ }
+ // Either the column is not dictionary encoded, or it was using dictionary encoding but
+ // a new data page has switched back to use plain encoding. For both cases we should
+ // return plain vector.
+ currentVector = cometVector;
+ return currentVector;
+ } else if (dictionary == null) {
+ // There is dictionary from native side but the Java side dictionary hasn't been
+ // initialized yet.
+ Dictionary arrowDictionary = dictionaryProvider.lookup(dictionaryEncoding.getId());
+ CometPlainVector dictionaryVector =
+ new CometPlainVector(arrowDictionary.getVector(), useDecimal128);
+ dictionary = new CometDictionary(dictionaryVector);
+ }
+
+ currentVector =
+ new CometDictionaryVector(cometVector, dictionary, dictionaryProvider, useDecimal128);
+ return currentVector;
+ }
+ }
+
+ protected void readPage() {
+ DataPage page = pageReader.readPage();
+ if (page == null) {
+ throw new RuntimeException("overreading: returned DataPage is null");
+ }
+ ;
+ int pageValueCount = page.getValueCount();
+ page.accept(
+ new DataPage.Visitor() {
+ @Override
+ public Void visit(DataPageV1 dataPageV1) {
+ LOG.debug("data page encoding = {}", dataPageV1.getValueEncoding());
+ if (dataPageV1.getDlEncoding() != Encoding.RLE
+ && descriptor.getMaxDefinitionLevel() != 0) {
+ throw new UnsupportedOperationException(
+ "Unsupported encoding: " + dataPageV1.getDlEncoding());
+ }
+ if (!isValidValueEncoding(dataPageV1.getValueEncoding())) {
+ throw new UnsupportedOperationException(
+ "Unsupported value encoding: " + dataPageV1.getValueEncoding());
+ }
+ try {
+ boolean useDirectBuffer =
+ (Boolean) CometConf.COMET_PARQUET_ENABLE_DIRECT_BUFFER().get();
+ if (useDirectBuffer) {
+ ByteBuffer buffer = dataPageV1.getBytes().toByteBuffer();
+ Native.setPageBufferV1(
+ nativeHandle, pageValueCount, buffer, dataPageV1.getValueEncoding().ordinal());
+ } else {
+ byte[] array = dataPageV1.getBytes().toByteArray();
+ Native.setPageV1(
+ nativeHandle, pageValueCount, array, dataPageV1.getValueEncoding().ordinal());
+ }
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ return null;
+ }
+
+ @Override
+ public Void visit(DataPageV2 dataPageV2) {
+ if (!isValidValueEncoding(dataPageV2.getDataEncoding())) {
+ throw new UnsupportedOperationException(
+ "Unsupported encoding: " + dataPageV2.getDataEncoding());
+ }
+ try {
+ Native.setPageV2(
+ nativeHandle,
+ pageValueCount,
+ dataPageV2.getDefinitionLevels().toByteArray(),
+ dataPageV2.getRepetitionLevels().toByteArray(),
+ dataPageV2.getData().toByteArray(),
+ dataPageV2.getDataEncoding().ordinal());
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ return null;
+ }
+ });
+ }
+
+ @SuppressWarnings("deprecation")
+ private boolean isValidValueEncoding(Encoding encoding) {
+ switch (encoding) {
+ case PLAIN:
+ case RLE_DICTIONARY:
+ case PLAIN_DICTIONARY:
+ return true;
+ default:
+ return false;
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/CometInputFile.java b/common/src/main/java/org/apache/comet/parquet/CometInputFile.java
new file mode 100644
index 000000000..eb54d1a72
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/CometInputFile.java
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataInputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.FutureDataInputStreamBuilder;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.util.VersionInfo;
+import org.apache.parquet.hadoop.util.HadoopStreams;
+import org.apache.parquet.io.InputFile;
+import org.apache.parquet.io.SeekableInputStream;
+
+/**
+ * A Parquet {@link InputFile} implementation that's similar to {@link
+ * org.apache.parquet.hadoop.util.HadoopInputFile}, but with optimizations introduced in Hadoop 3.x,
+ * for S3 specifically.
+ */
+public class CometInputFile implements InputFile {
+ private static final String MAJOR_MINOR_REGEX = "^(\\d+)\\.(\\d+)(\\..*)?$";
+ private static final Pattern VERSION_MATCHER = Pattern.compile(MAJOR_MINOR_REGEX);
+
+ private final FileSystem fs;
+ private final FileStatus stat;
+ private final Configuration conf;
+
+ public static CometInputFile fromPath(Path path, Configuration conf) throws IOException {
+ FileSystem fs = path.getFileSystem(conf);
+ return new CometInputFile(fs, fs.getFileStatus(path), conf);
+ }
+
+ private CometInputFile(FileSystem fs, FileStatus stat, Configuration conf) {
+ this.fs = fs;
+ this.stat = stat;
+ this.conf = conf;
+ }
+
+ @Override
+ public long getLength() {
+ return stat.getLen();
+ }
+
+ public Configuration getConf() {
+ return this.conf;
+ }
+
+ public FileSystem getFileSystem() {
+ return this.fs;
+ }
+
+ public Path getPath() {
+ return stat.getPath();
+ }
+
+ @Override
+ public SeekableInputStream newStream() throws IOException {
+ FSDataInputStream stream;
+ try {
+ if (isAtLeastHadoop33()) {
+ // If Hadoop version is >= 3.3.x, we'll use the 'openFile' API which can save a
+ // HEAD request from cloud storages like S3
+ FutureDataInputStreamBuilder inputStreamBuilder =
+ fs.openFile(stat.getPath()).withFileStatus(stat);
+
+ if (stat.getPath().toString().startsWith("s3a")) {
+ // Switch to random S3 input policy so that we don't do sequential read on the entire
+ // S3 object. By default, the policy is normal which does sequential read until a back
+ // seek happens, which in our case will never happen.
+ inputStreamBuilder =
+ inputStreamBuilder.opt("fs.s3a.experimental.input.fadvise", "random");
+ }
+ stream = inputStreamBuilder.build().get();
+ } else {
+ stream = fs.open(stat.getPath());
+ }
+ } catch (Exception e) {
+ throw new IOException("Error when opening file " + stat.getPath(), e);
+ }
+ return HadoopStreams.wrap(stream);
+ }
+
+ public SeekableInputStream newStream(long offset, long length) throws IOException {
+ try {
+ FSDataInputStream stream;
+ if (isAtLeastHadoop33()) {
+ FutureDataInputStreamBuilder inputStreamBuilder =
+ fs.openFile(stat.getPath()).withFileStatus(stat);
+
+ if (stat.getPath().toString().startsWith("s3a")) {
+ // Switch to random S3 input policy so that we don't do sequential read on the entire
+ // S3 object. By default, the policy is normal which does sequential read until a back
+ // seek happens, which in our case will never happen.
+ //
+ // Also set read ahead length equal to the column chunk length so we don't have to open
+ // multiple S3 http connections.
+ inputStreamBuilder =
+ inputStreamBuilder
+ .opt("fs.s3a.experimental.input.fadvise", "random")
+ .opt("fs.s3a.readahead.range", Long.toString(length));
+ }
+
+ stream = inputStreamBuilder.build().get();
+ } else {
+ stream = fs.open(stat.getPath());
+ }
+ return HadoopStreams.wrap(stream);
+ } catch (Exception e) {
+ throw new IOException(
+ "Error when opening file " + stat.getPath() + ", offset=" + offset + ", length=" + length,
+ e);
+ }
+ }
+
+ @Override
+ public String toString() {
+ return stat.getPath().toString();
+ }
+
+ private static boolean isAtLeastHadoop33() {
+ String version = VersionInfo.getVersion();
+ return CometInputFile.isAtLeastHadoop33(version);
+ }
+
+ static boolean isAtLeastHadoop33(String version) {
+ Matcher matcher = VERSION_MATCHER.matcher(version);
+ if (matcher.matches()) {
+ if (matcher.group(1).equals("3")) {
+ int minorVersion = Integer.parseInt(matcher.group(2));
+ return minorVersion >= 3;
+ }
+ }
+ return false;
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/ConstantColumnReader.java b/common/src/main/java/org/apache/comet/parquet/ConstantColumnReader.java
new file mode 100644
index 000000000..8de2376f9
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/ConstantColumnReader.java
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.math.BigInteger;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.types.*;
+import org.apache.spark.unsafe.types.UTF8String;
+
+import org.apache.comet.shims.ShimResolveDefaultColumns;
+
+/**
+ * A column reader that always return constant vectors. Used for reading partition columns, for
+ * instance.
+ */
+public class ConstantColumnReader extends MetadataColumnReader {
+ /** Whether all the values in this constant column are nulls */
+ private boolean isNull;
+
+ /** The constant value in the format of Object that are used to initialize this column reader. */
+ private Object value;
+
+ public ConstantColumnReader(StructField field, int batchSize, boolean useDecimal128) {
+ this(field.dataType(), TypeUtil.convertToParquet(field), batchSize, useDecimal128);
+ this.value = ShimResolveDefaultColumns.getExistenceDefaultValue(field);
+ init(value);
+ }
+
+ public ConstantColumnReader(
+ StructField field, int batchSize, InternalRow values, int index, boolean useDecimal128) {
+ this(field.dataType(), TypeUtil.convertToParquet(field), batchSize, useDecimal128);
+ init(values, index);
+ }
+
+ public ConstantColumnReader(
+ DataType type, ColumnDescriptor descriptor, Object value, boolean useDecimal128) {
+ super(type, descriptor, useDecimal128);
+ this.value = value;
+ }
+
+ ConstantColumnReader(
+ DataType type, ColumnDescriptor descriptor, int batchSize, boolean useDecimal128) {
+ super(type, descriptor, useDecimal128);
+ this.batchSize = batchSize;
+ initNative();
+ }
+
+ @Override
+ public void setBatchSize(int batchSize) {
+ super.setBatchSize(batchSize);
+ init(value);
+ }
+
+ @Override
+ public void readBatch(int total) {
+ super.readBatch(total);
+ if (isNull) setNumNulls(total);
+ }
+
+ private void init(InternalRow values, int index) {
+ Object value = values.get(index, type);
+ init(value);
+ }
+
+ private void init(Object value) {
+ if (value == null) {
+ Native.setNull(nativeHandle);
+ isNull = true;
+ } else if (type == DataTypes.BooleanType) {
+ Native.setBoolean(nativeHandle, (boolean) value);
+ } else if (type == DataTypes.ByteType) {
+ Native.setByte(nativeHandle, (byte) value);
+ } else if (type == DataTypes.ShortType) {
+ Native.setShort(nativeHandle, (short) value);
+ } else if (type == DataTypes.IntegerType) {
+ Native.setInt(nativeHandle, (int) value);
+ } else if (type == DataTypes.LongType) {
+ Native.setLong(nativeHandle, (long) value);
+ } else if (type == DataTypes.FloatType) {
+ Native.setFloat(nativeHandle, (float) value);
+ } else if (type == DataTypes.DoubleType) {
+ Native.setDouble(nativeHandle, (double) value);
+ } else if (type == DataTypes.BinaryType) {
+ Native.setBinary(nativeHandle, (byte[]) value);
+ } else if (type == DataTypes.StringType) {
+ Native.setBinary(nativeHandle, ((UTF8String) value).getBytes());
+ } else if (type == DataTypes.DateType) {
+ Native.setInt(nativeHandle, (int) value);
+ } else if (type == DataTypes.TimestampType || type == TimestampNTZType$.MODULE$) {
+ Native.setLong(nativeHandle, (long) value);
+ } else if (type instanceof DecimalType) {
+ DecimalType dt = (DecimalType) type;
+ Decimal d = (Decimal) value;
+ if (!useDecimal128 && dt.precision() <= Decimal.MAX_INT_DIGITS()) {
+ Native.setInt(nativeHandle, ((int) d.toUnscaledLong()));
+ } else if (!useDecimal128 && dt.precision() <= Decimal.MAX_LONG_DIGITS()) {
+ Native.setLong(nativeHandle, d.toUnscaledLong());
+ } else {
+ final BigInteger integer = d.toJavaBigDecimal().unscaledValue();
+ byte[] bytes = integer.toByteArray();
+ Native.setDecimal(nativeHandle, bytes);
+ }
+ } else {
+ throw new UnsupportedOperationException("Unsupported Spark type: " + type);
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/DictionaryPageReader.java b/common/src/main/java/org/apache/comet/parquet/DictionaryPageReader.java
new file mode 100644
index 000000000..b4b5a8fc9
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/DictionaryPageReader.java
@@ -0,0 +1,190 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.parquet.ParquetReadOptions;
+import org.apache.parquet.bytes.BytesInput;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.page.DictionaryPage;
+import org.apache.parquet.column.page.DictionaryPageReadStore;
+import org.apache.parquet.compression.CompressionCodecFactory;
+import org.apache.parquet.crypto.AesCipher;
+import org.apache.parquet.crypto.InternalColumnDecryptionSetup;
+import org.apache.parquet.crypto.InternalFileDecryptor;
+import org.apache.parquet.crypto.ModuleCipherFactory;
+import org.apache.parquet.format.BlockCipher;
+import org.apache.parquet.format.DictionaryPageHeader;
+import org.apache.parquet.format.PageHeader;
+import org.apache.parquet.format.Util;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.io.ParquetDecodingException;
+import org.apache.parquet.io.SeekableInputStream;
+
+public class DictionaryPageReader implements DictionaryPageReadStore {
+ private final Map> cache;
+ private final InternalFileDecryptor fileDecryptor;
+ private final SeekableInputStream inputStream;
+ private final ParquetReadOptions options;
+ private final Map columns;
+
+ DictionaryPageReader(
+ BlockMetaData block,
+ InternalFileDecryptor fileDecryptor,
+ SeekableInputStream inputStream,
+ ParquetReadOptions options) {
+ this.columns = new HashMap<>();
+ this.cache = new ConcurrentHashMap<>();
+ this.fileDecryptor = fileDecryptor;
+ this.inputStream = inputStream;
+ this.options = options;
+
+ for (ColumnChunkMetaData column : block.getColumns()) {
+ columns.put(column.getPath().toDotString(), column);
+ }
+ }
+
+ @Override
+ public DictionaryPage readDictionaryPage(ColumnDescriptor descriptor) {
+ String dotPath = String.join(".", descriptor.getPath());
+ ColumnChunkMetaData column = columns.get(dotPath);
+
+ if (column == null) {
+ throw new ParquetDecodingException("Failed to load dictionary, unknown column: " + dotPath);
+ }
+
+ return cache
+ .computeIfAbsent(
+ dotPath,
+ key -> {
+ try {
+ final DictionaryPage dict =
+ column.hasDictionaryPage() ? readDictionary(column) : null;
+
+ // Copy the dictionary to ensure it can be reused if it is returned
+ // more than once. This can happen when a DictionaryFilter has two or
+ // more predicates for the same column. Cache misses as well.
+ return (dict != null) ? Optional.of(reusableCopy(dict)) : Optional.empty();
+ } catch (IOException e) {
+ throw new ParquetDecodingException("Failed to read dictionary", e);
+ }
+ })
+ .orElse(null);
+ }
+
+ DictionaryPage readDictionary(ColumnChunkMetaData meta) throws IOException {
+ if (!meta.hasDictionaryPage()) {
+ return null;
+ }
+
+ if (inputStream.getPos() != meta.getStartingPos()) {
+ inputStream.seek(meta.getStartingPos());
+ }
+
+ boolean encryptedColumn = false;
+ InternalColumnDecryptionSetup columnDecryptionSetup = null;
+ byte[] dictionaryPageAAD = null;
+ BlockCipher.Decryptor pageDecryptor = null;
+ if (null != fileDecryptor && !fileDecryptor.plaintextFile()) {
+ columnDecryptionSetup = fileDecryptor.getColumnSetup(meta.getPath());
+ if (columnDecryptionSetup.isEncrypted()) {
+ encryptedColumn = true;
+ }
+ }
+
+ PageHeader pageHeader;
+ if (!encryptedColumn) {
+ pageHeader = Util.readPageHeader(inputStream);
+ } else {
+ byte[] dictionaryPageHeaderAAD =
+ AesCipher.createModuleAAD(
+ fileDecryptor.getFileAAD(),
+ ModuleCipherFactory.ModuleType.DictionaryPageHeader,
+ meta.getRowGroupOrdinal(),
+ columnDecryptionSetup.getOrdinal(),
+ -1);
+ pageHeader =
+ Util.readPageHeader(
+ inputStream, columnDecryptionSetup.getMetaDataDecryptor(), dictionaryPageHeaderAAD);
+ dictionaryPageAAD =
+ AesCipher.createModuleAAD(
+ fileDecryptor.getFileAAD(),
+ ModuleCipherFactory.ModuleType.DictionaryPage,
+ meta.getRowGroupOrdinal(),
+ columnDecryptionSetup.getOrdinal(),
+ -1);
+ pageDecryptor = columnDecryptionSetup.getDataDecryptor();
+ }
+
+ if (!pageHeader.isSetDictionary_page_header()) {
+ return null;
+ }
+
+ DictionaryPage compressedPage =
+ readCompressedDictionary(pageHeader, inputStream, pageDecryptor, dictionaryPageAAD);
+ CompressionCodecFactory.BytesInputDecompressor decompressor =
+ options.getCodecFactory().getDecompressor(meta.getCodec());
+
+ return new DictionaryPage(
+ decompressor.decompress(compressedPage.getBytes(), compressedPage.getUncompressedSize()),
+ compressedPage.getDictionarySize(),
+ compressedPage.getEncoding());
+ }
+
+ private DictionaryPage readCompressedDictionary(
+ PageHeader pageHeader,
+ SeekableInputStream fin,
+ BlockCipher.Decryptor pageDecryptor,
+ byte[] dictionaryPageAAD)
+ throws IOException {
+ DictionaryPageHeader dictHeader = pageHeader.getDictionary_page_header();
+
+ int uncompressedPageSize = pageHeader.getUncompressed_page_size();
+ int compressedPageSize = pageHeader.getCompressed_page_size();
+
+ byte[] dictPageBytes = new byte[compressedPageSize];
+ fin.readFully(dictPageBytes);
+
+ BytesInput bin = BytesInput.from(dictPageBytes);
+
+ if (null != pageDecryptor) {
+ bin = BytesInput.from(pageDecryptor.decrypt(bin.toByteArray(), dictionaryPageAAD));
+ }
+
+ return new DictionaryPage(
+ bin,
+ uncompressedPageSize,
+ dictHeader.getNum_values(),
+ org.apache.parquet.column.Encoding.valueOf(dictHeader.getEncoding().name()));
+ }
+
+ private static DictionaryPage reusableCopy(DictionaryPage dict) throws IOException {
+ return new DictionaryPage(
+ BytesInput.from(dict.getBytes().toByteArray()),
+ dict.getDictionarySize(),
+ dict.getEncoding());
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/FileReader.java b/common/src/main/java/org/apache/comet/parquet/FileReader.java
new file mode 100644
index 000000000..eddaf3f1a
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/FileReader.java
@@ -0,0 +1,1151 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.io.InputStream;
+import java.lang.reflect.Method;
+import java.net.URI;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.zip.CRC32;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.parquet.ParquetReadOptions;
+import org.apache.parquet.Preconditions;
+import org.apache.parquet.bytes.ByteBufferInputStream;
+import org.apache.parquet.bytes.BytesInput;
+import org.apache.parquet.bytes.BytesUtils;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.page.DataPage;
+import org.apache.parquet.column.page.DataPageV1;
+import org.apache.parquet.column.page.DataPageV2;
+import org.apache.parquet.column.page.DictionaryPage;
+import org.apache.parquet.column.page.PageReadStore;
+import org.apache.parquet.compression.CompressionCodecFactory;
+import org.apache.parquet.crypto.AesCipher;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.InternalColumnDecryptionSetup;
+import org.apache.parquet.crypto.InternalFileDecryptor;
+import org.apache.parquet.crypto.ModuleCipherFactory;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+import org.apache.parquet.filter2.compat.FilterCompat;
+import org.apache.parquet.format.BlockCipher;
+import org.apache.parquet.format.DataPageHeader;
+import org.apache.parquet.format.DataPageHeaderV2;
+import org.apache.parquet.format.DictionaryPageHeader;
+import org.apache.parquet.format.FileCryptoMetaData;
+import org.apache.parquet.format.PageHeader;
+import org.apache.parquet.format.Util;
+import org.apache.parquet.format.converter.ParquetMetadataConverter;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.hadoop.metadata.FileMetaData;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.apache.parquet.hadoop.util.counters.BenchmarkCounter;
+import org.apache.parquet.internal.column.columnindex.OffsetIndex;
+import org.apache.parquet.internal.filter2.columnindex.ColumnIndexFilter;
+import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore;
+import org.apache.parquet.internal.filter2.columnindex.RowRanges;
+import org.apache.parquet.io.InputFile;
+import org.apache.parquet.io.ParquetDecodingException;
+import org.apache.parquet.io.SeekableInputStream;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.spark.sql.execution.metric.SQLMetric;
+
+import static org.apache.parquet.hadoop.ParquetFileWriter.EFMAGIC;
+import static org.apache.parquet.hadoop.ParquetFileWriter.MAGIC;
+
+import static org.apache.comet.parquet.RowGroupFilter.FilterLevel.BLOOMFILTER;
+import static org.apache.comet.parquet.RowGroupFilter.FilterLevel.DICTIONARY;
+import static org.apache.comet.parquet.RowGroupFilter.FilterLevel.STATISTICS;
+
+/**
+ * A Parquet file reader. Mostly followed {@code ParquetFileReader} in {@code parquet-mr}, but with
+ * customizations & optimizations for Comet.
+ */
+public class FileReader implements Closeable {
+ private static final Logger LOG = LoggerFactory.getLogger(FileReader.class);
+
+ private final ParquetMetadataConverter converter;
+ protected final SeekableInputStream f;
+ private final InputFile file;
+ private final Map metrics;
+ private final Map paths = new HashMap<>();
+ private final FileMetaData fileMetaData; // may be null
+ private final List blocks;
+ private final List blockIndexStores;
+ private final List blockRowRanges;
+ private final CRC32 crc;
+ private final ParquetMetadata footer;
+
+ /**
+ * Read configurations come from two options: - options: these are options defined & specified
+ * from 'parquet-mr' library - cometOptions: these are Comet-specific options, for the features
+ * introduced in Comet's Parquet implementation
+ */
+ private final ParquetReadOptions options;
+
+ private final ReadOptions cometOptions;
+
+ private int currentBlock = 0;
+ private RowGroupReader currentRowGroup = null;
+ private InternalFileDecryptor fileDecryptor;
+
+ public FileReader(InputFile file, ParquetReadOptions options, ReadOptions cometOptions)
+ throws IOException {
+ this(file, null, options, cometOptions, null);
+ }
+
+ public FileReader(
+ InputFile file,
+ ParquetReadOptions options,
+ ReadOptions cometOptions,
+ Map metrics)
+ throws IOException {
+ this(file, null, options, cometOptions, metrics);
+ }
+
+ public FileReader(
+ InputFile file,
+ ParquetMetadata footer,
+ ParquetReadOptions options,
+ ReadOptions cometOptions,
+ Map metrics)
+ throws IOException {
+ this.converter = new ParquetMetadataConverter(options);
+ this.file = file;
+ this.f = file.newStream();
+ this.options = options;
+ this.cometOptions = cometOptions;
+ this.metrics = metrics;
+ if (footer == null) {
+ try {
+ footer = readFooter(file, options, f, converter);
+ } catch (Exception e) {
+ // In case that reading footer throws an exception in the constructor, the new stream
+ // should be closed. Otherwise, there's no way to close this outside.
+ f.close();
+ throw e;
+ }
+ }
+ this.footer = footer;
+ this.fileMetaData = footer.getFileMetaData();
+ this.fileDecryptor = fileMetaData.getFileDecryptor(); // must be called before filterRowGroups!
+ if (null != fileDecryptor && fileDecryptor.plaintextFile()) {
+ this.fileDecryptor = null; // Plaintext file. No need in decryptor
+ }
+
+ this.blocks = filterRowGroups(footer.getBlocks());
+ this.blockIndexStores = listWithNulls(this.blocks.size());
+ this.blockRowRanges = listWithNulls(this.blocks.size());
+ for (ColumnDescriptor col : footer.getFileMetaData().getSchema().getColumns()) {
+ paths.put(ColumnPath.get(col.getPath()), col);
+ }
+ this.crc = options.usePageChecksumVerification() ? new CRC32() : null;
+ }
+
+ /** Returns the footer of the Parquet file being read. */
+ public ParquetMetadata getFooter() {
+ return this.footer;
+ }
+
+ /** Returns the metadata of the Parquet file being read. */
+ public FileMetaData getFileMetaData() {
+ return this.fileMetaData;
+ }
+
+ /** Returns the input stream of the Parquet file being read. */
+ public SeekableInputStream getInputStream() {
+ return this.f;
+ }
+
+ /** Returns the Parquet options for reading the file. */
+ public ParquetReadOptions getOptions() {
+ return this.options;
+ }
+
+ /** Returns all the row groups of this reader (after applying row group filtering). */
+ public List getRowGroups() {
+ return blocks;
+ }
+
+ /** Sets the projected columns to be read later via {@link #readNextRowGroup()} */
+ public void setRequestedSchema(List projection) {
+ paths.clear();
+ for (ColumnDescriptor col : projection) {
+ paths.put(ColumnPath.get(col.getPath()), col);
+ }
+ }
+
+ /**
+ * Gets the total number of records across all row groups (after applying row group filtering).
+ */
+ public long getRecordCount() {
+ long total = 0;
+ for (BlockMetaData block : blocks) {
+ total += block.getRowCount();
+ }
+ return total;
+ }
+
+ /**
+ * Gets the total number of records across all row groups (after applying both row group filtering
+ * and page-level column index filtering).
+ */
+ public long getFilteredRecordCount() {
+ if (!options.useColumnIndexFilter()
+ || !FilterCompat.isFilteringRequired(options.getRecordFilter())) {
+ return getRecordCount();
+ }
+ long total = 0;
+ for (int i = 0, n = blocks.size(); i < n; ++i) {
+ total += getRowRanges(i).rowCount();
+ }
+ return total;
+ }
+
+ /** Skips the next row group. Returns false if there's no row group to skip. Otherwise, true. */
+ public boolean skipNextRowGroup() {
+ return advanceToNextBlock();
+ }
+
+ /**
+ * Returns the next row group to read (after applying row group filtering), or null if there's no
+ * more row group.
+ */
+ public PageReadStore readNextRowGroup() throws IOException {
+ if (currentBlock == blocks.size()) {
+ return null;
+ }
+ BlockMetaData block = blocks.get(currentBlock);
+ if (block.getRowCount() == 0) {
+ throw new RuntimeException("Illegal row group of 0 rows");
+ }
+ this.currentRowGroup = new RowGroupReader(block.getRowCount());
+ // prepare the list of consecutive parts to read them in one scan
+ List allParts = new ArrayList<>();
+ ConsecutivePartList currentParts = null;
+ for (ColumnChunkMetaData mc : block.getColumns()) {
+ ColumnPath pathKey = mc.getPath();
+ ColumnDescriptor columnDescriptor = paths.get(pathKey);
+ if (columnDescriptor != null) {
+ BenchmarkCounter.incrementTotalBytes(mc.getTotalSize());
+ long startingPos = mc.getStartingPos();
+ boolean mergeRanges = cometOptions.isIOMergeRangesEnabled();
+ int mergeRangeDelta = cometOptions.getIOMergeRangesDelta();
+
+ // start a new list if -
+ // it is the first part or
+ // the part is consecutive or
+ // the part is not consecutive but within the merge range
+ if (currentParts == null
+ || (!mergeRanges && currentParts.endPos() != startingPos)
+ || (mergeRanges && startingPos - currentParts.endPos() > mergeRangeDelta)) {
+ currentParts = new ConsecutivePartList(startingPos);
+ allParts.add(currentParts);
+ }
+ // if we are in a consecutive part list and there is a gap in between the parts,
+ // we treat the gap as a skippable chunk
+ long delta = startingPos - currentParts.endPos();
+ if (mergeRanges && delta > 0 && delta <= mergeRangeDelta) {
+ // add a chunk that will be skipped because it has no column descriptor
+ currentParts.addChunk(new ChunkDescriptor(null, null, startingPos, delta));
+ }
+ currentParts.addChunk(
+ new ChunkDescriptor(columnDescriptor, mc, startingPos, mc.getTotalSize()));
+ }
+ }
+ // actually read all the chunks
+ return readChunks(block, allParts, new ChunkListBuilder());
+ }
+
+ /**
+ * Returns the next row group to read (after applying both row group filtering and page level
+ * column index filtering), or null if there's no more row group.
+ */
+ public PageReadStore readNextFilteredRowGroup() throws IOException {
+ if (currentBlock == blocks.size()) {
+ return null;
+ }
+ if (!options.useColumnIndexFilter()
+ || !FilterCompat.isFilteringRequired(options.getRecordFilter())) {
+ return readNextRowGroup();
+ }
+ BlockMetaData block = blocks.get(currentBlock);
+ if (block.getRowCount() == 0) {
+ throw new RuntimeException("Illegal row group of 0 rows");
+ }
+ ColumnIndexStore ciStore = getColumnIndexReader(currentBlock);
+ RowRanges rowRanges = getRowRanges(currentBlock);
+ long rowCount = rowRanges.rowCount();
+ if (rowCount == 0) {
+ // There are no matching rows -> skipping this row-group
+ advanceToNextBlock();
+ return readNextFilteredRowGroup();
+ }
+ if (rowCount == block.getRowCount()) {
+ // All rows are matching -> fall back to the non-filtering path
+ return readNextRowGroup();
+ }
+
+ this.currentRowGroup = new RowGroupReader(rowRanges);
+ // prepare the list of consecutive parts to read them in one scan
+ ChunkListBuilder builder = new ChunkListBuilder();
+ List allParts = new ArrayList<>();
+ ConsecutivePartList currentParts = null;
+ for (ColumnChunkMetaData mc : block.getColumns()) {
+ ColumnPath pathKey = mc.getPath();
+ ColumnDescriptor columnDescriptor = paths.get(pathKey);
+ if (columnDescriptor != null) {
+ OffsetIndex offsetIndex = ciStore.getOffsetIndex(mc.getPath());
+ IndexFilter indexFilter = new IndexFilter(rowRanges, offsetIndex, block.getRowCount());
+ OffsetIndex filteredOffsetIndex = indexFilter.filterOffsetIndex();
+ for (IndexFilter.OffsetRange range :
+ indexFilter.calculateOffsetRanges(filteredOffsetIndex, mc)) {
+ BenchmarkCounter.incrementTotalBytes(range.length);
+ long startingPos = range.offset;
+ // first part or not consecutive => new list
+ if (currentParts == null || currentParts.endPos() != startingPos) {
+ currentParts = new ConsecutivePartList(startingPos);
+ allParts.add(currentParts);
+ }
+ ChunkDescriptor chunkDescriptor =
+ new ChunkDescriptor(columnDescriptor, mc, startingPos, range.length);
+ currentParts.addChunk(chunkDescriptor);
+ builder.setOffsetIndex(chunkDescriptor, filteredOffsetIndex);
+ }
+ }
+ }
+ // actually read all the chunks
+ return readChunks(block, allParts, builder);
+ }
+
+ // Visible for testing
+ ColumnIndexReader getColumnIndexReader(int blockIndex) {
+ ColumnIndexReader ciStore = blockIndexStores.get(blockIndex);
+ if (ciStore == null) {
+ ciStore = ColumnIndexReader.create(blocks.get(blockIndex), paths.keySet(), fileDecryptor, f);
+ blockIndexStores.set(blockIndex, ciStore);
+ }
+ return ciStore;
+ }
+
+ private PageReadStore readChunks(
+ BlockMetaData block, List allParts, ChunkListBuilder builder)
+ throws IOException {
+ for (ConsecutivePartList consecutiveChunks : allParts) {
+ if (shouldReadParallel()) {
+ consecutiveChunks.readAllParallel(builder);
+ } else {
+ consecutiveChunks.readAll(f, builder);
+ }
+ }
+ for (Chunk chunk : builder.build()) {
+ readChunkPages(chunk, block);
+ }
+
+ advanceToNextBlock();
+
+ return currentRowGroup;
+ }
+
+ private boolean shouldReadParallel() {
+ if (file instanceof CometInputFile) {
+ URI uri = ((CometInputFile) file).getPath().toUri();
+ return shouldReadParallel(cometOptions, uri.getScheme());
+ }
+
+ return false;
+ }
+
+ static boolean shouldReadParallel(ReadOptions options, String scheme) {
+ return options.isParallelIOEnabled() && shouldReadParallelForScheme(scheme);
+ }
+
+ private static boolean shouldReadParallelForScheme(String scheme) {
+ if (scheme == null) {
+ return false;
+ }
+
+ switch (scheme) {
+ case "s3a":
+ // Only enable parallel read for S3, so far.
+ return true;
+ default:
+ return false;
+ }
+ }
+
+ private void readChunkPages(Chunk chunk, BlockMetaData block) throws IOException {
+ if (fileDecryptor == null || fileDecryptor.plaintextFile()) {
+ currentRowGroup.addColumn(chunk.descriptor.col, chunk.readAllPages());
+ return;
+ }
+ // Encrypted file
+ ColumnPath columnPath = ColumnPath.get(chunk.descriptor.col.getPath());
+ InternalColumnDecryptionSetup columnDecryptionSetup = fileDecryptor.getColumnSetup(columnPath);
+ if (!columnDecryptionSetup.isEncrypted()) { // plaintext column
+ currentRowGroup.addColumn(chunk.descriptor.col, chunk.readAllPages());
+ } else { // encrypted column
+ currentRowGroup.addColumn(
+ chunk.descriptor.col,
+ chunk.readAllPages(
+ columnDecryptionSetup.getMetaDataDecryptor(),
+ columnDecryptionSetup.getDataDecryptor(),
+ fileDecryptor.getFileAAD(),
+ block.getOrdinal(),
+ columnDecryptionSetup.getOrdinal()));
+ }
+ }
+
+ private boolean advanceToNextBlock() {
+ if (currentBlock == blocks.size()) {
+ return false;
+ }
+ // update the current block and instantiate a dictionary reader for it
+ ++currentBlock;
+ return true;
+ }
+
+ public long[] getRowIndices() {
+ long[] rowIndices = new long[blocks.size() * 2];
+ for (int i = 0, n = blocks.size(); i < n; i++) {
+ BlockMetaData block = blocks.get(i);
+ rowIndices[i * 2] = getRowIndexOffset(block);
+ rowIndices[i * 2 + 1] = block.getRowCount();
+ }
+ return rowIndices;
+ }
+
+ // Uses reflection to get row index offset from a Parquet block metadata.
+ //
+ // The reason reflection is used here is that some Spark versions still depend on a
+ // Parquet version where the method `getRowIndexOffset` is not public.
+ private long getRowIndexOffset(BlockMetaData metaData) {
+ try {
+ Method method = BlockMetaData.class.getMethod("getRowIndexOffset");
+ method.setAccessible(true);
+ return (long) method.invoke(metaData);
+ } catch (Exception e) {
+ throw new RuntimeException("Error when calling getRowIndexOffset", e);
+ }
+ }
+
+ private RowRanges getRowRanges(int blockIndex) {
+ Preconditions.checkState(
+ FilterCompat.isFilteringRequired(options.getRecordFilter()),
+ "Should not be invoked if filter is null or NOOP");
+ RowRanges rowRanges = blockRowRanges.get(blockIndex);
+ if (rowRanges == null) {
+ rowRanges =
+ ColumnIndexFilter.calculateRowRanges(
+ options.getRecordFilter(),
+ getColumnIndexReader(blockIndex),
+ paths.keySet(),
+ blocks.get(blockIndex).getRowCount());
+ blockRowRanges.set(blockIndex, rowRanges);
+ }
+ return rowRanges;
+ }
+
+ private static ParquetMetadata readFooter(
+ InputFile file,
+ ParquetReadOptions options,
+ SeekableInputStream f,
+ ParquetMetadataConverter converter)
+ throws IOException {
+ long fileLen = file.getLength();
+ String filePath = file.toString();
+ LOG.debug("File length {}", fileLen);
+
+ int FOOTER_LENGTH_SIZE = 4;
+
+ // MAGIC + data + footer + footerIndex + MAGIC
+ if (fileLen < MAGIC.length + FOOTER_LENGTH_SIZE + MAGIC.length) {
+ throw new RuntimeException(
+ filePath + " is not a Parquet file (length is too low: " + fileLen + ")");
+ }
+
+ // Read footer length and magic string - with a single seek
+ byte[] magic = new byte[MAGIC.length];
+ long fileMetadataLengthIndex = fileLen - magic.length - FOOTER_LENGTH_SIZE;
+ LOG.debug("reading footer index at {}", fileMetadataLengthIndex);
+ f.seek(fileMetadataLengthIndex);
+ int fileMetadataLength = BytesUtils.readIntLittleEndian(f);
+ f.readFully(magic);
+
+ boolean encryptedFooterMode;
+ if (Arrays.equals(MAGIC, magic)) {
+ encryptedFooterMode = false;
+ } else if (Arrays.equals(EFMAGIC, magic)) {
+ encryptedFooterMode = true;
+ } else {
+ throw new RuntimeException(
+ filePath
+ + " is not a Parquet file. Expected magic number "
+ + "at tail, but found "
+ + Arrays.toString(magic));
+ }
+
+ long fileMetadataIndex = fileMetadataLengthIndex - fileMetadataLength;
+ LOG.debug("read footer length: {}, footer index: {}", fileMetadataLength, fileMetadataIndex);
+ if (fileMetadataIndex < magic.length || fileMetadataIndex >= fileMetadataLengthIndex) {
+ throw new RuntimeException(
+ "corrupted file: the footer index is not within the file: " + fileMetadataIndex);
+ }
+ f.seek(fileMetadataIndex);
+
+ FileDecryptionProperties fileDecryptionProperties = options.getDecryptionProperties();
+ InternalFileDecryptor fileDecryptor = null;
+ if (null != fileDecryptionProperties) {
+ fileDecryptor = new InternalFileDecryptor(fileDecryptionProperties);
+ }
+
+ // Read all the footer bytes in one time to avoid multiple read operations,
+ // since it can be pretty time consuming for a single read operation in HDFS.
+ byte[] footerBytes = new byte[fileMetadataLength];
+ f.readFully(footerBytes);
+ ByteBuffer footerBytesBuffer = ByteBuffer.wrap(footerBytes);
+ LOG.debug("Finished to read all footer bytes.");
+ InputStream footerBytesStream = ByteBufferInputStream.wrap(footerBytesBuffer);
+
+ // Regular file, or encrypted file with plaintext footer
+ if (!encryptedFooterMode) {
+ return converter.readParquetMetadata(
+ footerBytesStream, options.getMetadataFilter(), fileDecryptor, false, fileMetadataLength);
+ }
+
+ // Encrypted file with encrypted footer
+ if (fileDecryptor == null) {
+ throw new ParquetCryptoRuntimeException(
+ "Trying to read file with encrypted footer. " + "No keys available");
+ }
+ FileCryptoMetaData fileCryptoMetaData = Util.readFileCryptoMetaData(footerBytesStream);
+ fileDecryptor.setFileCryptoMetaData(
+ fileCryptoMetaData.getEncryption_algorithm(), true, fileCryptoMetaData.getKey_metadata());
+ // footer length is required only for signed plaintext footers
+ return converter.readParquetMetadata(
+ footerBytesStream, options.getMetadataFilter(), fileDecryptor, true, 0);
+ }
+
+ private List filterRowGroups(List blocks) {
+ FilterCompat.Filter recordFilter = options.getRecordFilter();
+ if (FilterCompat.isFilteringRequired(recordFilter)) {
+ // set up data filters based on configured levels
+ List levels = new ArrayList<>();
+
+ if (options.useStatsFilter()) {
+ levels.add(STATISTICS);
+ }
+
+ if (options.useDictionaryFilter()) {
+ levels.add(DICTIONARY);
+ }
+
+ if (options.useBloomFilter()) {
+ levels.add(BLOOMFILTER);
+ }
+ return RowGroupFilter.filterRowGroups(levels, recordFilter, blocks, this);
+ }
+
+ return blocks;
+ }
+
+ private static List listWithNulls(int size) {
+ return Stream.generate(() -> (T) null).limit(size).collect(Collectors.toList());
+ }
+
+ public void closeStream() throws IOException {
+ if (f != null) {
+ f.close();
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
+ try {
+ if (f != null) {
+ f.close();
+ }
+ } finally {
+ options.getCodecFactory().release();
+ }
+ }
+
+ /**
+ * Builder to concatenate the buffers of the discontinuous parts for the same column. These parts
+ * are generated as a result of the column-index based filtering when some pages might be skipped
+ * at reading.
+ */
+ private class ChunkListBuilder {
+ private class ChunkData {
+ final List buffers = new ArrayList<>();
+ OffsetIndex offsetIndex;
+ }
+
+ private final Map map = new HashMap<>();
+
+ void add(ChunkDescriptor descriptor, List buffers) {
+ ChunkListBuilder.ChunkData data = map.get(descriptor);
+ if (data == null) {
+ data = new ChunkData();
+ map.put(descriptor, data);
+ }
+ data.buffers.addAll(buffers);
+ }
+
+ void setOffsetIndex(ChunkDescriptor descriptor, OffsetIndex offsetIndex) {
+ ChunkData data = map.get(descriptor);
+ if (data == null) {
+ data = new ChunkData();
+ map.put(descriptor, data);
+ }
+ data.offsetIndex = offsetIndex;
+ }
+
+ List build() {
+ List chunks = new ArrayList<>();
+ for (Map.Entry entry : map.entrySet()) {
+ ChunkDescriptor descriptor = entry.getKey();
+ ChunkData data = entry.getValue();
+ chunks.add(new Chunk(descriptor, data.buffers, data.offsetIndex));
+ }
+ return chunks;
+ }
+ }
+
+ /** The data for a column chunk */
+ private class Chunk {
+ private final ChunkDescriptor descriptor;
+ private final ByteBufferInputStream stream;
+ final OffsetIndex offsetIndex;
+
+ /**
+ * @param descriptor descriptor for the chunk
+ * @param buffers ByteBuffers that contain the chunk
+ * @param offsetIndex the offset index for this column; might be null
+ */
+ Chunk(ChunkDescriptor descriptor, List buffers, OffsetIndex offsetIndex) {
+ this.descriptor = descriptor;
+ this.stream = ByteBufferInputStream.wrap(buffers);
+ this.offsetIndex = offsetIndex;
+ }
+
+ protected PageHeader readPageHeader(BlockCipher.Decryptor blockDecryptor, byte[] pageHeaderAAD)
+ throws IOException {
+ return Util.readPageHeader(stream, blockDecryptor, pageHeaderAAD);
+ }
+
+ /**
+ * Calculate checksum of input bytes, throw decoding exception if it does not match the provided
+ * reference crc
+ */
+ private void verifyCrc(int referenceCrc, byte[] bytes, String exceptionMsg) {
+ crc.reset();
+ crc.update(bytes);
+ if (crc.getValue() != ((long) referenceCrc & 0xffffffffL)) {
+ throw new ParquetDecodingException(exceptionMsg);
+ }
+ }
+
+ private ColumnPageReader readAllPages() throws IOException {
+ return readAllPages(null, null, null, -1, -1);
+ }
+
+ private ColumnPageReader readAllPages(
+ BlockCipher.Decryptor headerBlockDecryptor,
+ BlockCipher.Decryptor pageBlockDecryptor,
+ byte[] aadPrefix,
+ int rowGroupOrdinal,
+ int columnOrdinal)
+ throws IOException {
+ List pagesInChunk = new ArrayList<>();
+ DictionaryPage dictionaryPage = null;
+ PrimitiveType type =
+ fileMetaData.getSchema().getType(descriptor.col.getPath()).asPrimitiveType();
+
+ long valuesCountReadSoFar = 0;
+ int dataPageCountReadSoFar = 0;
+ byte[] dataPageHeaderAAD = null;
+ if (null != headerBlockDecryptor) {
+ dataPageHeaderAAD =
+ AesCipher.createModuleAAD(
+ aadPrefix,
+ ModuleCipherFactory.ModuleType.DataPageHeader,
+ rowGroupOrdinal,
+ columnOrdinal,
+ getPageOrdinal(dataPageCountReadSoFar));
+ }
+ while (hasMorePages(valuesCountReadSoFar, dataPageCountReadSoFar)) {
+ byte[] pageHeaderAAD = dataPageHeaderAAD;
+ if (null != headerBlockDecryptor) {
+ // Important: this verifies file integrity (makes sure dictionary page had not been
+ // removed)
+ if (null == dictionaryPage && descriptor.metadata.hasDictionaryPage()) {
+ pageHeaderAAD =
+ AesCipher.createModuleAAD(
+ aadPrefix,
+ ModuleCipherFactory.ModuleType.DictionaryPageHeader,
+ rowGroupOrdinal,
+ columnOrdinal,
+ -1);
+ } else {
+ int pageOrdinal = getPageOrdinal(dataPageCountReadSoFar);
+ AesCipher.quickUpdatePageAAD(dataPageHeaderAAD, pageOrdinal);
+ }
+ }
+
+ PageHeader pageHeader = readPageHeader(headerBlockDecryptor, pageHeaderAAD);
+ int uncompressedPageSize = pageHeader.getUncompressed_page_size();
+ int compressedPageSize = pageHeader.getCompressed_page_size();
+ final BytesInput pageBytes;
+ switch (pageHeader.type) {
+ case DICTIONARY_PAGE:
+ // there is only one dictionary page per column chunk
+ if (dictionaryPage != null) {
+ throw new ParquetDecodingException(
+ "more than one dictionary page in column " + descriptor.col);
+ }
+ pageBytes = this.readAsBytesInput(compressedPageSize);
+ if (options.usePageChecksumVerification() && pageHeader.isSetCrc()) {
+ verifyCrc(
+ pageHeader.getCrc(),
+ pageBytes.toByteArray(),
+ "could not verify dictionary page integrity, CRC checksum verification failed");
+ }
+ DictionaryPageHeader dicHeader = pageHeader.getDictionary_page_header();
+ dictionaryPage =
+ new DictionaryPage(
+ pageBytes,
+ uncompressedPageSize,
+ dicHeader.getNum_values(),
+ converter.getEncoding(dicHeader.getEncoding()));
+ // Copy crc to new page, used for testing
+ if (pageHeader.isSetCrc()) {
+ dictionaryPage.setCrc(pageHeader.getCrc());
+ }
+ break;
+
+ case DATA_PAGE:
+ DataPageHeader dataHeaderV1 = pageHeader.getData_page_header();
+ pageBytes = this.readAsBytesInput(compressedPageSize);
+ if (options.usePageChecksumVerification() && pageHeader.isSetCrc()) {
+ verifyCrc(
+ pageHeader.getCrc(),
+ pageBytes.toByteArray(),
+ "could not verify page integrity, CRC checksum verification failed");
+ }
+ DataPageV1 dataPageV1 =
+ new DataPageV1(
+ pageBytes,
+ dataHeaderV1.getNum_values(),
+ uncompressedPageSize,
+ converter.fromParquetStatistics(
+ getFileMetaData().getCreatedBy(), dataHeaderV1.getStatistics(), type),
+ converter.getEncoding(dataHeaderV1.getRepetition_level_encoding()),
+ converter.getEncoding(dataHeaderV1.getDefinition_level_encoding()),
+ converter.getEncoding(dataHeaderV1.getEncoding()));
+ // Copy crc to new page, used for testing
+ if (pageHeader.isSetCrc()) {
+ dataPageV1.setCrc(pageHeader.getCrc());
+ }
+ pagesInChunk.add(dataPageV1);
+ valuesCountReadSoFar += dataHeaderV1.getNum_values();
+ ++dataPageCountReadSoFar;
+ break;
+
+ case DATA_PAGE_V2:
+ DataPageHeaderV2 dataHeaderV2 = pageHeader.getData_page_header_v2();
+ int dataSize =
+ compressedPageSize
+ - dataHeaderV2.getRepetition_levels_byte_length()
+ - dataHeaderV2.getDefinition_levels_byte_length();
+ pagesInChunk.add(
+ new DataPageV2(
+ dataHeaderV2.getNum_rows(),
+ dataHeaderV2.getNum_nulls(),
+ dataHeaderV2.getNum_values(),
+ this.readAsBytesInput(dataHeaderV2.getRepetition_levels_byte_length()),
+ this.readAsBytesInput(dataHeaderV2.getDefinition_levels_byte_length()),
+ converter.getEncoding(dataHeaderV2.getEncoding()),
+ this.readAsBytesInput(dataSize),
+ uncompressedPageSize,
+ converter.fromParquetStatistics(
+ getFileMetaData().getCreatedBy(), dataHeaderV2.getStatistics(), type),
+ dataHeaderV2.isIs_compressed()));
+ valuesCountReadSoFar += dataHeaderV2.getNum_values();
+ ++dataPageCountReadSoFar;
+ break;
+
+ default:
+ LOG.debug(
+ "skipping page of type {} of size {}", pageHeader.getType(), compressedPageSize);
+ stream.skipFully(compressedPageSize);
+ break;
+ }
+ }
+ if (offsetIndex == null && valuesCountReadSoFar != descriptor.metadata.getValueCount()) {
+ // Would be nice to have a CorruptParquetFileException or something as a subclass?
+ throw new IOException(
+ "Expected "
+ + descriptor.metadata.getValueCount()
+ + " values in column chunk at "
+ + file
+ + " offset "
+ + descriptor.metadata.getFirstDataPageOffset()
+ + " but got "
+ + valuesCountReadSoFar
+ + " values instead over "
+ + pagesInChunk.size()
+ + " pages ending at file offset "
+ + (descriptor.fileOffset + stream.position()));
+ }
+ CompressionCodecFactory.BytesInputDecompressor decompressor =
+ options.getCodecFactory().getDecompressor(descriptor.metadata.getCodec());
+ return new ColumnPageReader(
+ decompressor,
+ pagesInChunk,
+ dictionaryPage,
+ offsetIndex,
+ blocks.get(currentBlock).getRowCount(),
+ pageBlockDecryptor,
+ aadPrefix,
+ rowGroupOrdinal,
+ columnOrdinal);
+ }
+
+ private boolean hasMorePages(long valuesCountReadSoFar, int dataPageCountReadSoFar) {
+ return offsetIndex == null
+ ? valuesCountReadSoFar < descriptor.metadata.getValueCount()
+ : dataPageCountReadSoFar < offsetIndex.getPageCount();
+ }
+
+ private int getPageOrdinal(int dataPageCountReadSoFar) {
+ if (null == offsetIndex) {
+ return dataPageCountReadSoFar;
+ }
+
+ return offsetIndex.getPageOrdinal(dataPageCountReadSoFar);
+ }
+
+ /**
+ * @param size the size of the page
+ * @return the page
+ * @throws IOException if there is an error while reading from the file stream
+ */
+ public BytesInput readAsBytesInput(int size) throws IOException {
+ return BytesInput.from(stream.sliceBuffers(size));
+ }
+ }
+
+ /**
+ * Describes a list of consecutive parts to be read at once. A consecutive part may contain whole
+ * column chunks or only parts of them (some pages).
+ */
+ private class ConsecutivePartList {
+ private final long offset;
+ private final List chunks = new ArrayList<>();
+ private long length;
+ private final SQLMetric fileReadTimeMetric;
+ private final SQLMetric fileReadSizeMetric;
+ private final SQLMetric readThroughput;
+
+ /**
+ * @param offset where the first chunk starts
+ */
+ ConsecutivePartList(long offset) {
+ if (metrics != null) {
+ this.fileReadTimeMetric = metrics.get("ParquetInputFileReadTime");
+ this.fileReadSizeMetric = metrics.get("ParquetInputFileReadSize");
+ this.readThroughput = metrics.get("ParquetInputFileReadThroughput");
+ } else {
+ this.fileReadTimeMetric = null;
+ this.fileReadSizeMetric = null;
+ this.readThroughput = null;
+ }
+ this.offset = offset;
+ }
+
+ /**
+ * Adds a chunk to the list. It must be consecutive to the previous chunk.
+ *
+ * @param descriptor a chunk descriptor
+ */
+ public void addChunk(ChunkDescriptor descriptor) {
+ chunks.add(descriptor);
+ length += descriptor.size;
+ }
+
+ private List allocateReadBuffers() {
+ int fullAllocations = Math.toIntExact(length / options.getMaxAllocationSize());
+ int lastAllocationSize = Math.toIntExact(length % options.getMaxAllocationSize());
+
+ int numAllocations = fullAllocations + (lastAllocationSize > 0 ? 1 : 0);
+ List buffers = new ArrayList<>(numAllocations);
+
+ for (int i = 0; i < fullAllocations; i += 1) {
+ buffers.add(options.getAllocator().allocate(options.getMaxAllocationSize()));
+ }
+
+ if (lastAllocationSize > 0) {
+ buffers.add(options.getAllocator().allocate(lastAllocationSize));
+ }
+ return buffers;
+ }
+
+ /**
+ * @param f file to read the chunks from
+ * @param builder used to build chunk list to read the pages for the different columns
+ * @throws IOException if there is an error while reading from the stream
+ */
+ public void readAll(SeekableInputStream f, ChunkListBuilder builder) throws IOException {
+ f.seek(offset);
+
+ List buffers = allocateReadBuffers();
+ long startNs = System.nanoTime();
+
+ for (ByteBuffer buffer : buffers) {
+ f.readFully(buffer);
+ buffer.flip();
+ }
+ setReadMetrics(startNs);
+
+ // report in a counter the data we just scanned
+ BenchmarkCounter.incrementBytesRead(length);
+ ByteBufferInputStream stream = ByteBufferInputStream.wrap(buffers);
+ for (int i = 0; i < chunks.size(); i++) {
+ ChunkDescriptor descriptor = chunks.get(i);
+ if (descriptor.col != null) {
+ builder.add(descriptor, stream.sliceBuffers(descriptor.size));
+ } else {
+ stream.skipFully(descriptor.size);
+ }
+ }
+ }
+
+ /**
+ * Api to read a consecutive range from the Parquet file in parallel. This is identical to
+ * {@link #readAll(SeekableInputStream, ChunkListBuilder) readAll}, except that the consecutive
+ * range is split into multiple smaller ranges and read in parallel. The parallelism can be set
+ * by specifying the threadpool size via {@link
+ * ReadOptions.Builder#withParallelIOThreadPoolSize(int)}.
+ *
+ * @param builder used to build chunk list to read the pages for the different columns
+ * @throws IOException if there is an error while reading from the stream
+ */
+ public void readAllParallel(ChunkListBuilder builder) throws IOException {
+
+ List buffers = allocateReadBuffers();
+ long startNs = System.nanoTime();
+
+ int nThreads = cometOptions.parallelIOThreadPoolSize();
+ ExecutorService threadPool = CometFileReaderThreadPool.getOrCreateThreadPool(nThreads);
+ List> futures = new ArrayList<>();
+
+ long currentOffset = this.offset;
+ int buffersPerThread = buffers.size() / nThreads;
+ int remaining = buffers.size() % nThreads;
+ // offset in input file each thread seeks to before beginning read
+ long[] offsets = new long[nThreads];
+ // index of buffer where each thread will start writing data
+ int[] bufferIndexes = new int[nThreads];
+ // number of buffers for each thread to fill
+ int[] numBuffers = new int[nThreads];
+
+ int bufferNum = 0;
+ for (int i = 0; i < nThreads; i++) {
+ int nBuffers = 0;
+ offsets[i] = currentOffset;
+ bufferIndexes[i] = bufferNum;
+ nBuffers = buffersPerThread;
+ for (int j = 0; j < buffersPerThread; j++) {
+ currentOffset += buffers.get(bufferNum).capacity();
+ bufferNum++;
+ }
+ if (remaining > 0) {
+ remaining--;
+ currentOffset += buffers.get(bufferNum).capacity();
+ bufferNum++;
+ nBuffers++;
+ }
+ numBuffers[i] = nBuffers;
+ }
+ for (int n = 0; n < nThreads; n++) {
+ int threadIndex = n;
+ long pos = offsets[threadIndex];
+ int bufferIndex = bufferIndexes[threadIndex];
+ int nBuffers = numBuffers[threadIndex];
+ if (nBuffers == 0) {
+ continue;
+ }
+
+ // Find the total number of bytes to read for the current thread
+ long tmp = 0;
+ for (int i = 0; i < nBuffers; i++) {
+ int bufNo = bufferIndex + i;
+ if (bufNo >= buffers.size()) break;
+ tmp += buffers.get(bufNo).capacity();
+ }
+ final long length = tmp;
+
+ futures.add(
+ threadPool.submit(
+ () -> {
+ SeekableInputStream inputStream = null;
+ try {
+ if (file instanceof CometInputFile) {
+ inputStream = (((CometInputFile) file).newStream(pos, length));
+ } else {
+ inputStream = file.newStream();
+ }
+
+ inputStream.seek(pos);
+ long curPos = pos;
+ for (int i = 0; i < nBuffers; i++) {
+ int bufNo = bufferIndex + i;
+ if (bufNo >= buffers.size()) {
+ break;
+ }
+ ByteBuffer buffer = buffers.get(bufNo);
+ LOG.debug(
+ "Thread: {} Offset: {} Buffer: {} Size: {}",
+ threadIndex,
+ curPos,
+ bufNo,
+ buffer.capacity());
+ curPos += buffer.capacity();
+ inputStream.readFully(buffer);
+ buffer.flip();
+ } // for
+ } finally {
+ if (inputStream != null) {
+ inputStream.close();
+ }
+ }
+
+ return null;
+ }));
+ }
+
+ for (Future future : futures) {
+ try {
+ future.get();
+ } catch (InterruptedException | ExecutionException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ setReadMetrics(startNs);
+
+ ByteBufferInputStream stream;
+ stream = ByteBufferInputStream.wrap(buffers);
+ // report in a counter the data we just scanned
+ BenchmarkCounter.incrementBytesRead(length);
+ for (int i = 0; i < chunks.size(); i++) {
+ ChunkDescriptor descriptor = chunks.get(i);
+ if (descriptor.col != null) {
+ builder.add(descriptor, stream.sliceBuffers(descriptor.size));
+ } else {
+ stream.skipFully(descriptor.size);
+ }
+ }
+ }
+
+ private void setReadMetrics(long startNs) {
+ long totalFileReadTimeNs = System.nanoTime() - startNs;
+ double sizeInMb = ((double) length) / (1024 * 1024);
+ double timeInSec = ((double) totalFileReadTimeNs) / 1000_0000_0000L;
+ double throughput = sizeInMb / timeInSec;
+ LOG.debug(
+ "Comet: File Read stats: Length: {} MB, Time: {} secs, throughput: {} MB/sec ",
+ sizeInMb,
+ timeInSec,
+ throughput);
+ if (fileReadTimeMetric != null) {
+ fileReadTimeMetric.add(totalFileReadTimeNs);
+ }
+ if (fileReadSizeMetric != null) {
+ fileReadSizeMetric.add(length);
+ }
+ if (readThroughput != null) {
+ readThroughput.set(throughput);
+ }
+ }
+
+ /**
+ * @return the position following the last byte of these chunks
+ */
+ public long endPos() {
+ return offset + length;
+ }
+ }
+
+ /** Information needed to read a column chunk or a part of it. */
+ private static class ChunkDescriptor {
+ private final ColumnDescriptor col;
+ private final ColumnChunkMetaData metadata;
+ private final long fileOffset;
+ private final long size;
+
+ /**
+ * @param col column this chunk is part of
+ * @param metadata metadata for the column
+ * @param fileOffset offset in the file where this chunk starts
+ * @param size size of the chunk
+ */
+ ChunkDescriptor(
+ ColumnDescriptor col, ColumnChunkMetaData metadata, long fileOffset, long size) {
+ this.col = col;
+ this.metadata = metadata;
+ this.fileOffset = fileOffset;
+ this.size = size;
+ }
+
+ @Override
+ public int hashCode() {
+ return col.hashCode();
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj) {
+ return true;
+ } else if (obj instanceof ChunkDescriptor) {
+ return col.equals(((ChunkDescriptor) obj).col);
+ } else {
+ return false;
+ }
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/FooterReader.java b/common/src/main/java/org/apache/comet/parquet/FooterReader.java
new file mode 100644
index 000000000..092eb442c
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/FooterReader.java
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.HadoopReadOptions;
+import org.apache.parquet.ParquetReadOptions;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.apache.spark.sql.execution.datasources.PartitionedFile;
+
+/**
+ * Copied from Spark's `ParquetFooterReader` in order to avoid shading issue around Parquet.
+ *
+ *
`FooterReader` is a util class which encapsulates the helper methods of reading parquet file
+ * footer.
+ */
+public class FooterReader {
+ public static ParquetMetadata readFooter(Configuration configuration, PartitionedFile file)
+ throws IOException, URISyntaxException {
+ long start = file.start();
+ long length = file.length();
+ Path filePath = new Path(new URI(file.filePath().toString()));
+ CometInputFile inputFile = CometInputFile.fromPath(filePath, configuration);
+ ParquetReadOptions readOptions =
+ HadoopReadOptions.builder(inputFile.getConf(), inputFile.getPath())
+ .withRange(start, start + length)
+ .build();
+ ReadOptions cometReadOptions = ReadOptions.builder(configuration).build();
+ // Use try-with-resources to ensure fd is closed.
+ try (FileReader fileReader = new FileReader(inputFile, readOptions, cometReadOptions)) {
+ return fileReader.getFooter();
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/IndexFilter.java b/common/src/main/java/org/apache/comet/parquet/IndexFilter.java
new file mode 100644
index 000000000..afa5687ca
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/IndexFilter.java
@@ -0,0 +1,145 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.internal.column.columnindex.OffsetIndex;
+import org.apache.parquet.internal.filter2.columnindex.RowRanges;
+
+public class IndexFilter {
+ private final RowRanges rowRanges;
+ private final OffsetIndex offsetIndex;
+ private final long totalRowCount;
+
+ public IndexFilter(RowRanges rowRanges, OffsetIndex offsetIndex, long totalRowCount) {
+ this.rowRanges = rowRanges;
+ this.offsetIndex = offsetIndex;
+ this.totalRowCount = totalRowCount;
+ }
+
+ OffsetIndex filterOffsetIndex() {
+ List indexMap = new ArrayList<>();
+ for (int i = 0, n = offsetIndex.getPageCount(); i < n; ++i) {
+ long from = offsetIndex.getFirstRowIndex(i);
+ if (rowRanges.isOverlapping(from, offsetIndex.getLastRowIndex(i, totalRowCount))) {
+ indexMap.add(i);
+ }
+ }
+
+ int[] indexArray = new int[indexMap.size()];
+ for (int i = 0; i < indexArray.length; i++) {
+ indexArray[i] = indexMap.get(i);
+ }
+ return new FilteredOffsetIndex(offsetIndex, indexArray);
+ }
+
+ List calculateOffsetRanges(OffsetIndex filteredOffsetIndex, ColumnChunkMetaData cm) {
+ List ranges = new ArrayList<>();
+ long firstPageOffset = offsetIndex.getOffset(0);
+ int n = filteredOffsetIndex.getPageCount();
+
+ if (n > 0) {
+ OffsetRange currentRange = null;
+
+ // Add a range for the dictionary page if required
+ long rowGroupOffset = cm.getStartingPos();
+ if (rowGroupOffset < firstPageOffset) {
+ currentRange = new OffsetRange(rowGroupOffset, (int) (firstPageOffset - rowGroupOffset));
+ ranges.add(currentRange);
+ }
+
+ for (int i = 0; i < n; ++i) {
+ long offset = filteredOffsetIndex.getOffset(i);
+ int length = filteredOffsetIndex.getCompressedPageSize(i);
+ if (currentRange == null || !currentRange.extend(offset, length)) {
+ currentRange = new OffsetRange(offset, length);
+ ranges.add(currentRange);
+ }
+ }
+ }
+ return ranges;
+ }
+
+ private static class FilteredOffsetIndex implements OffsetIndex {
+ private final OffsetIndex offsetIndex;
+ private final int[] indexMap;
+
+ private FilteredOffsetIndex(OffsetIndex offsetIndex, int[] indexMap) {
+ this.offsetIndex = offsetIndex;
+ this.indexMap = indexMap;
+ }
+
+ @Override
+ public int getPageOrdinal(int pageIndex) {
+ return indexMap[pageIndex];
+ }
+
+ @Override
+ public int getPageCount() {
+ return indexMap.length;
+ }
+
+ @Override
+ public long getOffset(int pageIndex) {
+ return offsetIndex.getOffset(indexMap[pageIndex]);
+ }
+
+ @Override
+ public int getCompressedPageSize(int pageIndex) {
+ return offsetIndex.getCompressedPageSize(indexMap[pageIndex]);
+ }
+
+ @Override
+ public long getFirstRowIndex(int pageIndex) {
+ return offsetIndex.getFirstRowIndex(indexMap[pageIndex]);
+ }
+
+ @Override
+ public long getLastRowIndex(int pageIndex, long totalRowCount) {
+ int nextIndex = indexMap[pageIndex] + 1;
+ return (nextIndex >= offsetIndex.getPageCount()
+ ? totalRowCount
+ : offsetIndex.getFirstRowIndex(nextIndex))
+ - 1;
+ }
+ }
+
+ static class OffsetRange {
+ final long offset;
+ long length;
+
+ private OffsetRange(long offset, int length) {
+ this.offset = offset;
+ this.length = length;
+ }
+
+ private boolean extend(long offset, int length) {
+ if (this.offset + this.length == offset) {
+ this.length += length;
+ return true;
+ } else {
+ return false;
+ }
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/LazyColumnReader.java b/common/src/main/java/org/apache/comet/parquet/LazyColumnReader.java
new file mode 100644
index 000000000..a15d84192
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/LazyColumnReader.java
@@ -0,0 +1,186 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.io.IOException;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.page.PageReader;
+import org.apache.spark.sql.types.DataType;
+
+import org.apache.comet.vector.CometLazyVector;
+import org.apache.comet.vector.CometVector;
+
+public class LazyColumnReader extends ColumnReader {
+
+ // Remember the largest skipped index for sanity checking.
+ private int lastSkippedRowId = Integer.MAX_VALUE;
+
+ // Track whether the underlying page is drained.
+ private boolean isPageDrained = true;
+
+ // Leftover number of rows that did not skip in the previous batch.
+ private int numRowsToSkipFromPrevBatch;
+
+ // The lazy vector being updated.
+ private final CometLazyVector vector;
+
+ public LazyColumnReader(
+ DataType sparkReadType,
+ ColumnDescriptor descriptor,
+ int batchSize,
+ boolean useDecimal128,
+ boolean useLegacyDateTimestamp) {
+ super(sparkReadType, descriptor, batchSize, useDecimal128, useLegacyDateTimestamp);
+ this.batchSize = 0; // the batch size is set later in `readBatch`
+ this.vector = new CometLazyVector(sparkReadType, this, useDecimal128);
+ }
+
+ @Override
+ public void setPageReader(PageReader pageReader) throws IOException {
+ super.setPageReader(pageReader);
+ lastSkippedRowId = Integer.MAX_VALUE;
+ isPageDrained = true;
+ numRowsToSkipFromPrevBatch = 0;
+ currentNumValues = batchSize;
+ }
+
+ /**
+ * Lazily read a batch of 'total' rows for this column. The includes: 1) Skip any unused rows from
+ * the previous batch 2) Reset the native columnar batch 3) Reset tracking variables
+ *
+ * @param total the number of rows in the batch. MUST be <= the number of rows available in this
+ * column chunk.
+ */
+ @Override
+ public void readBatch(int total) {
+ // Before starting a new batch, take care of the remaining rows to skip from the previous batch.
+ tryPageSkip(batchSize);
+ numRowsToSkipFromPrevBatch += batchSize - currentNumValues;
+
+ // Now first reset the current columnar batch so that it can be used to fill in a new batch
+ // of values. Then, keep reading more data pages (via 'readBatch') until the current batch is
+ // full, or we have read 'total' number of values.
+ Native.resetBatch(nativeHandle);
+
+ batchSize = total;
+ currentNumValues = 0;
+ lastSkippedRowId = -1;
+ }
+
+ @Override
+ public CometVector currentBatch() {
+ return vector;
+ }
+
+ /** Read all rows up to the `batchSize`. Expects no rows are skipped so far. */
+ public void readAllBatch() {
+ // All rows should be read without any skips so far
+ assert (lastSkippedRowId == -1);
+
+ readBatch(batchSize - 1, 0);
+ }
+
+ /**
+ * Read at least up to `rowId`. It may read beyond `rowId` if enough rows available in the page.
+ * It may skip reading rows before `rowId`. In case `rowId` is already read, return immediately.
+ *
+ * @param rowId the row index in the batch to read.
+ * @return true if `rowId` is newly materialized, or false if `rowId` is already materialized.
+ */
+ public boolean materializeUpToIfNecessary(int rowId) {
+ // Not allowed reading rowId if it may have skipped previously.
+ assert (rowId > lastSkippedRowId);
+
+ // If `rowId` is already materialized, return immediately.
+ if (rowId < currentNumValues) return false;
+
+ int numRowsWholePageSkipped = tryPageSkip(rowId);
+ readBatch(rowId, numRowsWholePageSkipped);
+ return true;
+ }
+
+ /**
+ * Read up to `rowId` (inclusive). If the whole pages are skipped previously in `tryPageSkip()`,
+ * pad the number of whole page skipped rows with nulls to the underlying vector before reading.
+ *
+ * @param rowId the row index in the batch to read.
+ * @param numNullRowsToPad the number of nulls to pad before reading.
+ */
+ private void readBatch(int rowId, int numNullRowsToPad) {
+ if (numRowsToSkipFromPrevBatch > 0) {
+ // Reaches here only when starting a new batch and the page is previously drained
+ readPage();
+ isPageDrained = false;
+ Native.skipBatch(nativeHandle, numRowsToSkipFromPrevBatch, true);
+ numRowsToSkipFromPrevBatch = 0;
+ }
+ while (rowId >= currentNumValues) {
+ int numRowsToRead = batchSize - currentNumValues;
+ if (isPageDrained) {
+ readPage();
+ }
+ int[] array = Native.readBatch(nativeHandle, numRowsToRead, numNullRowsToPad);
+ int read = array[0];
+ isPageDrained = read < numRowsToRead;
+ currentNumValues += read;
+ currentNumNulls += array[1];
+ // No need to update numNullRowsToPad. numNullRowsToPad > 0 means there were whole page skips.
+ // That guarantees that the Native.readBatch can read up to rowId in the current page.
+ }
+ }
+
+ /**
+ * Try to skip until `rowId` (exclusive). If possible, it skips whole underlying pages without
+ * decompressing. In that case, it returns early at the page end, so that the next iteration can
+ * lazily decide to `readPage()` or `tryPageSkip()` again.
+ *
+ * @param rowId the row index in the batch that it tries to skip up until (exclusive).
+ * @return the number of rows that the whole page skips were applied.
+ */
+ private int tryPageSkip(int rowId) {
+ int total = rowId - currentNumValues;
+ int wholePageSkipped = 0;
+ if (total > 0) {
+ // First try to skip from the non-drained underlying page.
+ int skipped = isPageDrained ? 0 : Native.skipBatch(nativeHandle, total);
+ total -= skipped;
+ isPageDrained = total > 0;
+ if (isPageDrained) {
+ ColumnPageReader columnPageReader = (ColumnPageReader) pageReader;
+ // It is always `columnPageReader.getPageValueCount() > numRowsToSkipFromPriorBatch`
+ int pageValueCount = columnPageReader.getPageValueCount() - numRowsToSkipFromPrevBatch;
+ while (pageValueCount <= total) {
+ // skip the entire page if the next page is small enough
+ columnPageReader.skipPage();
+ numRowsToSkipFromPrevBatch = 0;
+ total -= pageValueCount;
+ wholePageSkipped += pageValueCount;
+ pageValueCount = columnPageReader.getPageValueCount();
+ }
+ }
+
+ currentNumValues += skipped + wholePageSkipped;
+ currentNumNulls += skipped;
+ lastSkippedRowId = currentNumValues - 1;
+ }
+ return wholePageSkipped;
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/MetadataColumnReader.java b/common/src/main/java/org/apache/comet/parquet/MetadataColumnReader.java
new file mode 100644
index 000000000..b8722ca78
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/MetadataColumnReader.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import org.apache.arrow.c.ArrowArray;
+import org.apache.arrow.c.ArrowSchema;
+import org.apache.arrow.c.Data;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.spark.sql.types.DataType;
+
+import org.apache.comet.vector.CometPlainVector;
+import org.apache.comet.vector.CometVector;
+
+/** A metadata column reader that can be extended by {@link RowIndexColumnReader} etc. */
+public class MetadataColumnReader extends AbstractColumnReader {
+ private final BufferAllocator allocator = new RootAllocator();
+ private CometVector vector;
+
+ public MetadataColumnReader(DataType type, ColumnDescriptor descriptor, boolean useDecimal128) {
+ // TODO: should we handle legacy dates & timestamps for metadata columns?
+ super(type, descriptor, useDecimal128, false);
+ }
+
+ @Override
+ public void setBatchSize(int batchSize) {
+ close();
+ super.setBatchSize(batchSize);
+ }
+
+ @Override
+ public void readBatch(int total) {
+ if (vector == null) {
+ long[] addresses = Native.currentBatch(nativeHandle);
+ try (ArrowArray array = ArrowArray.wrap(addresses[0]);
+ ArrowSchema schema = ArrowSchema.wrap(addresses[1])) {
+ FieldVector fieldVector = Data.importVector(allocator, array, schema, null);
+ vector = new CometPlainVector(fieldVector, useDecimal128);
+ }
+ }
+ vector.setNumValues(total);
+ }
+
+ void setNumNulls(int total) {
+ vector.setNumNulls(total);
+ }
+
+ @Override
+ public CometVector currentBatch() {
+ return vector;
+ }
+
+ @Override
+ public void close() {
+ if (vector != null) {
+ vector.close();
+ vector = null;
+ }
+ super.close();
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/Native.java b/common/src/main/java/org/apache/comet/parquet/Native.java
new file mode 100644
index 000000000..0887ae12f
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/Native.java
@@ -0,0 +1,235 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.nio.ByteBuffer;
+
+import org.apache.comet.NativeBase;
+
+public final class Native extends NativeBase {
+ public static int[] readBatch(long handle, int batchSize) {
+ return readBatch(handle, batchSize, 0);
+ }
+
+ public static int skipBatch(long handle, int batchSize) {
+ return skipBatch(handle, batchSize, false);
+ }
+
+ /** Native APIs * */
+
+ /**
+ * Creates a reader for a primitive Parquet column.
+ *
+ * @param physicalTypeId id for Parquet physical type
+ * @param logicalTypeId id for Parquet logical type
+ * @param expectedPhysicalTypeId id for Parquet physical type, converted from Spark read type.
+ * This is used for type promotion.
+ * @param path the path from the root schema to the column, derived from the method
+ * 'ColumnDescriptor#getPath()'.
+ * @param maxDl the maximum definition level of the primitive column
+ * @param maxRl the maximum repetition level of the primitive column
+ * @param bitWidth (only set when logical type is INT) the bit width for the integer type (INT8,
+ * INT16, INT32, etc)
+ * @param isSigned (only set when logical type is INT) whether it is signed or unsigned int.
+ * @param typeLength number of bytes required to store a value of the type, only set when the
+ * physical type is FIXED_LEN_BYTE_ARRAY, otherwise it's 0.
+ * @param precision (only set when logical type is DECIMAL) precision of the decimal type
+ * @param expectedPrecision (only set when logical type is DECIMAL) precision of the decimal type
+ * from Spark read schema. This is used for type promotion.
+ * @param scale (only set when logical type is DECIMAL) scale of the decimal type
+ * @param tu (only set when logical type is TIMESTAMP) unit for the timestamp
+ * @param isAdjustedUtc (only set when logical type is TIMESTAMP) whether the timestamp is
+ * adjusted to UTC or not
+ * @param batchSize the batch size for the columnar read
+ * @param useDecimal128 whether to always return 128 bit decimal regardless of precision
+ * @param useLegacyDateTimestampOrNTZ whether to read legacy dates/timestamps as it is
+ * @return a pointer to a native Parquet column reader created
+ */
+ public static native long initColumnReader(
+ int physicalTypeId,
+ int logicalTypeId,
+ int expectedPhysicalTypeId,
+ String[] path,
+ int maxDl,
+ int maxRl,
+ int bitWidth,
+ boolean isSigned,
+ int typeLength,
+ int precision,
+ int expectedPrecision,
+ int scale,
+ int tu,
+ boolean isAdjustedUtc,
+ int batchSize,
+ boolean useDecimal128,
+ boolean useLegacyDateTimestampOrNTZ);
+
+ /**
+ * Pass a Parquet dictionary page to the native column reader. Note this should only be called
+ * once per Parquet column chunk. Otherwise it'll panic.
+ *
+ * @param handle the handle to the native Parquet column reader
+ * @param dictionaryValueCount the number of values in this dictionary
+ * @param dictionaryData the actual dictionary page data, including repetition/definition levels
+ * as well as values
+ * @param encoding the encoding used by the dictionary
+ */
+ public static native void setDictionaryPage(
+ long handle, int dictionaryValueCount, byte[] dictionaryData, int encoding);
+
+ /**
+ * Passes a Parquet data page V1 to the native column reader.
+ *
+ * @param handle the handle to the native Parquet column reader
+ * @param pageValueCount the number of values in this data page
+ * @param pageData the actual page data, which should only contain PLAIN-encoded values.
+ * @param valueEncoding the encoding used by the values
+ */
+ public static native void setPageV1(
+ long handle, int pageValueCount, byte[] pageData, int valueEncoding);
+
+ /**
+ * Passes a Parquet data page V1 to the native column reader.
+ *
+ * @param handle the handle to the native Parquet column reader
+ * @param pageValueCount the number of values in this data page
+ * @param buffer the actual page data, represented by a DirectByteBuffer.
+ * @param valueEncoding the encoding used by the values
+ */
+ public static native void setPageBufferV1(
+ long handle, int pageValueCount, ByteBuffer buffer, int valueEncoding);
+
+ /**
+ * Passes a Parquet data page V2 to the native column reader.
+ *
+ * @param handle the handle to the native Parquet column reader
+ * @param pageValueCount the number of values in this data page
+ * @param defLevelData the data for definition levels
+ * @param repLevelData the data for repetition levels
+ * @param valueData the data for values
+ * @param valueEncoding the encoding used by the values
+ */
+ public static native void setPageV2(
+ long handle,
+ int pageValueCount,
+ byte[] defLevelData,
+ byte[] repLevelData,
+ byte[] valueData,
+ int valueEncoding);
+
+ /**
+ * Reset the current columnar batch. This will clear all the content of the batch as well as any
+ * internal state such as the current offset.
+ *
+ * @param handle the handle to the native Parquet column reader
+ */
+ public static native void resetBatch(long handle);
+
+ /**
+ * Reads at most 'batchSize' number of rows from the native Parquet column reader. Returns a tuple
+ * where the first element is the actual number of rows read (including both nulls and non-nulls),
+ * and the second element is the number of nulls read.
+ *
+ *
If the returned value is < 'batchSize' then it means the current page has been completely
+ * drained. In this case, the caller should call {@link Native#setPageV1} or {@link
+ * Native#setPageV2} before the next 'readBatch' call.
+ *
+ *
Note that the current page could also be drained if the returned value = 'batchSize', i.e.,
+ * the remaining number of rows in the page is exactly equal to 'batchSize'. In this case, the
+ * next 'readBatch' call will return 0 and the caller should call {@link Native#setPageV1} or
+ * {@link Native#setPageV2} next.
+ *
+ *
If `nullPadSize` > 0, it pads nulls into the underlying vector before the values will be
+ * read into.
+ *
+ * @param handle the handle to the native Parquet column reader
+ * @param batchSize the number of rows to be read
+ * @param nullPadSize the number of nulls to pad before reading.
+ * @return a tuple: (the actual number of rows read, the number of nulls read)
+ */
+ public static native int[] readBatch(long handle, int batchSize, int nullPadSize);
+
+ /**
+ * Skips at most 'batchSize' number of rows from the native Parquet column reader, and returns the
+ * actual number of rows skipped.
+ *
+ *
If the returned value is < 'batchSize' then it means the current page has been completely
+ * drained. In this case, the caller should call {@link Native#setPageV1} or {@link
+ * Native#setPageV2} before the next 'skipBatch' call.
+ *
+ *
Note that the current page could also be drained if the returned value = 'batchSize', i.e.,
+ * the remaining number of rows in the page is exactly equal to 'batchSize'. In this case, the
+ * next 'skipBatch' call will return 0 and the caller should call {@link Native#setPageV1} or
+ * {@link Native#setPageV2} next.
+ *
+ * @param handle the handle to the native Parquet column reader
+ * @param batchSize the number of rows to skip in the current page
+ * @param discard if true, discard read rows without padding nulls into the underlying vector
+ * @return the actual number of rows skipped
+ */
+ public static native int skipBatch(long handle, int batchSize, boolean discard);
+
+ /**
+ * Returns the current batch constructed via 'readBatch'
+ *
+ * @param handle the handle to the native Parquet column reader
+ * @return a long array with 2 elements, the first is the address to native Arrow array, and the
+ * second is the address to the Arrow schema.
+ */
+ public static native long[] currentBatch(long handle);
+
+ /** Set methods to set a constant value for the reader, so it'll return constant vectors */
+ public static native void setNull(long handle);
+
+ public static native void setBoolean(long handle, boolean value);
+
+ public static native void setByte(long handle, byte value);
+
+ public static native void setShort(long handle, short value);
+
+ public static native void setInt(long handle, int value);
+
+ public static native void setLong(long handle, long value);
+
+ public static native void setFloat(long handle, float value);
+
+ public static native void setDouble(long handle, double value);
+
+ public static native void setBinary(long handle, byte[] value);
+
+ /** Set decimal backed by FixedLengthByteArray */
+ public static native void setDecimal(long handle, byte[] value);
+
+ /** Set position of row index vector for Iceberg Metadata Column */
+ public static native void setPosition(long handle, long value, int size);
+
+ /** Set row index vector for Spark row index metadata column and return vector size */
+ public static native int setIndices(long handle, long offset, int size, long[] indices);
+
+ /** Set deleted info for Iceberg Metadata Column */
+ public static native void setIsDeleted(long handle, boolean[] isDeleted);
+
+ /**
+ * Closes the native Parquet column reader and releases all resources associated with it.
+ *
+ * @param handle the handle to the native Parquet column reader
+ */
+ public static native void closeColumnReader(long handle);
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/ReadOptions.java b/common/src/main/java/org/apache/comet/parquet/ReadOptions.java
new file mode 100644
index 000000000..6754443e6
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/ReadOptions.java
@@ -0,0 +1,197 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.spark.SparkEnv;
+import org.apache.spark.launcher.SparkLauncher;
+
+/**
+ * Comet specific Parquet related read options.
+ *
+ *
TODO: merge this with {@link org.apache.parquet.HadoopReadOptions} once PARQUET-2203 is done.
+ */
+public class ReadOptions {
+ private static final Logger LOG = LoggerFactory.getLogger(ReadOptions.class);
+ public static final String COMET_PARQUET_PARALLEL_IO_ENABLED =
+ "comet.parquet.read.parallel.io.enabled";
+ public static final boolean COMET_PARQUET_PARALLEL_IO_ENABLED_DEFAULT = true;
+
+ public static final String COMET_PARQUET_PARALLEL_IO_THREADS =
+ "comet.parquet.read.parallel.io.thread-pool.size";
+ public static final int COMET_PARQUET_PARALLEL_IO_THREADS_DEFAULT = 32;
+
+ public static final String COMET_IO_MERGE_RANGES = "comet.parquet.read.io.mergeRanges";
+ private static final boolean COMET_IO_MERGE_RANGES_DEFAULT = true;
+
+ public static final String COMET_IO_MERGE_RANGES_DELTA =
+ "comet.parquet.read.io.mergeRanges.delta";
+ private static final int COMET_IO_MERGE_RANGES_DELTA_DEFAULT = 1 << 23; // 8 MB
+
+ // Max number of concurrent tasks we expect. Used to autoconfigure S3 client connections
+ public static final int S3A_MAX_EXPECTED_PARALLELISM = 32;
+ // defined in hadoop-aws - org.apache.hadoop.fs.s3a.Constants.MAXIMUM_CONNECTIONS
+ public static final String S3A_MAXIMUM_CONNECTIONS = "fs.s3a.connection.maximum";
+ // default max connections in S3A - org.apache.hadoop.fs.s3a.Constants.DEFAULT_MAXIMUM_CONNECTIONS
+ public static final int S3A_DEFAULT_MAX_HTTP_CONNECTIONS = 96;
+
+ public static final String S3A_READAHEAD_RANGE = "fs.s3a.readahead.range";
+ // Default read ahead range in Hadoop is 64K; we increase it to 1 MB
+ public static final long COMET_DEFAULT_READAHEAD_RANGE = 1 * 1024 * 1024; // 1 MB
+
+ private final boolean parallelIOEnabled;
+ private final int parallelIOThreadPoolSize;
+ private final boolean ioMergeRanges;
+ private final int ioMergeRangesDelta;
+
+ ReadOptions(
+ boolean parallelIOEnabled,
+ int parallelIOThreadPoolSize,
+ boolean ioMergeRanges,
+ int ioMergeRangesDelta) {
+ this.parallelIOEnabled = parallelIOEnabled;
+ this.parallelIOThreadPoolSize = parallelIOThreadPoolSize;
+ this.ioMergeRanges = ioMergeRanges;
+ this.ioMergeRangesDelta = ioMergeRangesDelta;
+ }
+
+ public boolean isParallelIOEnabled() {
+ return this.parallelIOEnabled;
+ }
+
+ public int parallelIOThreadPoolSize() {
+ return this.parallelIOThreadPoolSize;
+ }
+
+ public boolean isIOMergeRangesEnabled() {
+ return ioMergeRanges;
+ }
+
+ public int getIOMergeRangesDelta() {
+ return ioMergeRangesDelta;
+ }
+
+ public static Builder builder(Configuration conf) {
+ return new Builder(conf);
+ }
+
+ public static Builder builder() {
+ return builder(new Configuration());
+ }
+
+ public static class Builder {
+ private final Configuration conf;
+
+ private boolean parallelIOEnabled;
+ private int parallelIOThreadPoolSize;
+ private boolean ioMergeRanges;
+ private int ioMergeRangesDelta;
+
+ /**
+ * Whether to enable Parquet parallel IO when reading row groups. If true, Parquet reader will
+ * use multiple threads to read multiple chunks of data from the current row group in parallel.
+ */
+ public Builder enableParallelIO(boolean b) {
+ this.parallelIOEnabled = b;
+ return this;
+ }
+
+ /**
+ * Specify the number of threads to be used in parallel IO.
+ *
+ *
Note: this will only be effective if parallel IO is enabled (e.g., via {@link
+ * #enableParallelIO(boolean)}).
+ */
+ public Builder withParallelIOThreadPoolSize(int numThreads) {
+ this.parallelIOThreadPoolSize = numThreads;
+ return this;
+ }
+
+ public Builder enableIOMergeRanges(boolean enableIOMergeRanges) {
+ this.ioMergeRanges = enableIOMergeRanges;
+ return this;
+ }
+
+ public Builder withIOMergeRangesDelta(int ioMergeRangesDelta) {
+ this.ioMergeRangesDelta = ioMergeRangesDelta;
+ return this;
+ }
+
+ public ReadOptions build() {
+ return new ReadOptions(
+ parallelIOEnabled, parallelIOThreadPoolSize, ioMergeRanges, ioMergeRangesDelta);
+ }
+
+ public Builder(Configuration conf) {
+ this.conf = conf;
+ this.parallelIOEnabled =
+ conf.getBoolean(
+ COMET_PARQUET_PARALLEL_IO_ENABLED, COMET_PARQUET_PARALLEL_IO_ENABLED_DEFAULT);
+ this.parallelIOThreadPoolSize =
+ conf.getInt(COMET_PARQUET_PARALLEL_IO_THREADS, COMET_PARQUET_PARALLEL_IO_THREADS_DEFAULT);
+ this.ioMergeRanges = conf.getBoolean(COMET_IO_MERGE_RANGES, COMET_IO_MERGE_RANGES_DEFAULT);
+ this.ioMergeRangesDelta =
+ conf.getInt(COMET_IO_MERGE_RANGES_DELTA, COMET_IO_MERGE_RANGES_DELTA_DEFAULT);
+ // override some S3 defaults
+ setS3Config();
+ }
+
+ // For paths to S3, if the s3 connection pool max is less than twice the product of
+ // parallel reader threads * number of cores, then increase the connection pool max
+ private void setS3Config() {
+ int s3ConnectionsMax = S3A_DEFAULT_MAX_HTTP_CONNECTIONS;
+ SparkEnv env = SparkEnv.get();
+ // Use a default number of cores in case we are using the FileReader outside the context
+ // of Spark.
+ int numExecutorCores = S3A_MAX_EXPECTED_PARALLELISM;
+ if (env != null) {
+ numExecutorCores = env.conf().getInt(SparkLauncher.EXECUTOR_CORES, numExecutorCores);
+ }
+ int parallelReaderThreads = this.parallelIOEnabled ? this.parallelIOThreadPoolSize : 1;
+ s3ConnectionsMax = Math.max(numExecutorCores * parallelReaderThreads * 2, s3ConnectionsMax);
+
+ setS3ConfIfGreater(conf, S3A_MAXIMUM_CONNECTIONS, s3ConnectionsMax);
+ setS3ConfIfGreater(conf, S3A_READAHEAD_RANGE, COMET_DEFAULT_READAHEAD_RANGE);
+ }
+
+ // Update the conf iff the new value is greater than the existing val
+ private void setS3ConfIfGreater(Configuration conf, String key, int newVal) {
+ int maxVal = newVal;
+ String curr = conf.get(key);
+ if (curr != null && !curr.isEmpty()) {
+ maxVal = Math.max(Integer.parseInt(curr), newVal);
+ }
+ LOG.info("File reader auto configured '{}={}'", key, maxVal);
+ conf.set(key, Integer.toString(maxVal));
+ }
+
+ // Update the conf iff the new value is greater than the existing val. This handles values that
+ // may have suffixes (K, M, G, T, P, E) indicating well known bytes size suffixes
+ private void setS3ConfIfGreater(Configuration conf, String key, long newVal) {
+ long maxVal = conf.getLongBytes(key, newVal);
+ maxVal = Math.max(maxVal, newVal);
+ LOG.info("File reader auto configured '{}={}'", key, maxVal);
+ conf.set(key, Long.toString(maxVal));
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/RowGroupFilter.java b/common/src/main/java/org/apache/comet/parquet/RowGroupFilter.java
new file mode 100644
index 000000000..40a61c05e
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/RowGroupFilter.java
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.parquet.filter2.compat.FilterCompat;
+import org.apache.parquet.filter2.compat.FilterCompat.Filter;
+import org.apache.parquet.filter2.compat.FilterCompat.NoOpFilter;
+import org.apache.parquet.filter2.compat.FilterCompat.Visitor;
+import org.apache.parquet.filter2.dictionarylevel.DictionaryFilter;
+import org.apache.parquet.filter2.predicate.FilterPredicate;
+import org.apache.parquet.filter2.predicate.SchemaCompatibilityValidator;
+import org.apache.parquet.filter2.statisticslevel.StatisticsFilter;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.schema.MessageType;
+
+public class RowGroupFilter implements Visitor> {
+ private final List blocks;
+ private final MessageType schema;
+ private final List levels;
+ private final FileReader reader;
+
+ public enum FilterLevel {
+ STATISTICS,
+ DICTIONARY,
+ BLOOMFILTER
+ }
+
+ public static List filterRowGroups(
+ List levels, Filter filter, List blocks, FileReader reader) {
+ return filter.accept(new RowGroupFilter(levels, blocks, reader));
+ }
+
+ private RowGroupFilter(List levels, List blocks, FileReader reader) {
+ this.levels = levels;
+ this.blocks = blocks;
+ this.reader = reader;
+ this.schema = reader.getFileMetaData().getSchema();
+ }
+
+ @Override
+ public List visit(FilterCompat.FilterPredicateCompat filterPredicateCompat) {
+ FilterPredicate filterPredicate = filterPredicateCompat.getFilterPredicate();
+
+ // check that the schema of the filter matches the schema of the file
+ SchemaCompatibilityValidator.validate(filterPredicate, schema);
+
+ List filteredBlocks = new ArrayList<>();
+
+ for (BlockMetaData block : blocks) {
+ boolean drop = false;
+
+ if (levels.contains(FilterLevel.STATISTICS)) {
+ drop = StatisticsFilter.canDrop(filterPredicate, block.getColumns());
+ }
+
+ if (!drop && levels.contains(FilterLevel.DICTIONARY)) {
+ drop =
+ DictionaryFilter.canDrop(
+ filterPredicate,
+ block.getColumns(),
+ new DictionaryPageReader(
+ block,
+ reader.getFileMetaData().getFileDecryptor(),
+ reader.getInputStream(),
+ reader.getOptions()));
+ }
+
+ if (!drop && levels.contains(FilterLevel.BLOOMFILTER)) {
+ drop =
+ filterPredicate.accept(
+ new BloomFilterReader(
+ block, reader.getFileMetaData().getFileDecryptor(), reader.getInputStream()));
+ }
+
+ if (!drop) {
+ filteredBlocks.add(block);
+ }
+ }
+
+ return filteredBlocks;
+ }
+
+ @Override
+ public List visit(
+ FilterCompat.UnboundRecordFilterCompat unboundRecordFilterCompat) {
+ return blocks;
+ }
+
+ @Override
+ public List visit(NoOpFilter noOpFilter) {
+ return blocks;
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/RowGroupReader.java b/common/src/main/java/org/apache/comet/parquet/RowGroupReader.java
new file mode 100644
index 000000000..d5d73b078
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/RowGroupReader.java
@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.PrimitiveIterator;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.page.PageReadStore;
+import org.apache.parquet.column.page.PageReader;
+import org.apache.parquet.internal.filter2.columnindex.RowRanges;
+
+public class RowGroupReader implements PageReadStore {
+ private final Map readers = new HashMap<>();
+ private final long rowCount;
+ private final RowRanges rowRanges;
+
+ public RowGroupReader(long rowCount) {
+ this.rowCount = rowCount;
+ this.rowRanges = null;
+ }
+
+ RowGroupReader(RowRanges rowRanges) {
+ this.rowRanges = rowRanges;
+ this.rowCount = rowRanges.rowCount();
+ }
+
+ @Override
+ public long getRowCount() {
+ return rowCount;
+ }
+
+ @Override
+ public PageReader getPageReader(ColumnDescriptor path) {
+ final PageReader pageReader = readers.get(path);
+ if (pageReader == null) {
+ throw new IllegalArgumentException(
+ path + " is not found: " + readers.keySet() + " " + rowCount);
+ }
+ return pageReader;
+ }
+
+ @Override
+ public Optional getRowIndexes() {
+ return rowRanges == null ? Optional.empty() : Optional.of(rowRanges.iterator());
+ }
+
+ void addColumn(ColumnDescriptor path, ColumnPageReader reader) {
+ if (readers.put(path, reader) != null) {
+ throw new IllegalStateException(path + " was already added");
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/RowIndexColumnReader.java b/common/src/main/java/org/apache/comet/parquet/RowIndexColumnReader.java
new file mode 100644
index 000000000..8448318db
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/RowIndexColumnReader.java
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import org.apache.spark.sql.types.*;
+
+/**
+ * A column reader that returns the row index vector. Used for reading row index metadata column for
+ * Spark 3.4+. The row index can be accessed by {@code _tmp_metadata_row_index} column.
+ */
+public class RowIndexColumnReader extends MetadataColumnReader {
+ /** The row indices that are used to initialize this column reader. */
+ private final long[] indices;
+
+ /** The current number of indices to skip reading from {@code indices}. */
+ private long offset;
+
+ public RowIndexColumnReader(StructField field, int batchSize, long[] indices) {
+ super(field.dataType(), TypeUtil.convertToParquet(field), false);
+ this.indices = indices;
+ setBatchSize(batchSize);
+ }
+
+ @Override
+ public void readBatch(int total) {
+ Native.resetBatch(nativeHandle);
+ int count = Native.setIndices(nativeHandle, offset, total, indices);
+ offset += count;
+
+ super.readBatch(count);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/SupportsComet.java b/common/src/main/java/org/apache/comet/parquet/SupportsComet.java
new file mode 100644
index 000000000..f330a7681
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/SupportsComet.java
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+public interface SupportsComet {
+ boolean isCometEnabled();
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/TypeUtil.java b/common/src/main/java/org/apache/comet/parquet/TypeUtil.java
new file mode 100644
index 000000000..340bf9823
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/TypeUtil.java
@@ -0,0 +1,279 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.util.Arrays;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.schema.LogicalTypeAnnotation;
+import org.apache.parquet.schema.LogicalTypeAnnotation.*;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.parquet.schema.Type;
+import org.apache.parquet.schema.Types;
+import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException;
+import org.apache.spark.sql.types.*;
+
+import org.apache.comet.CometConf;
+
+public class TypeUtil {
+
+ /** Converts the input Spark 'field' into a Parquet column descriptor. */
+ public static ColumnDescriptor convertToParquet(StructField field) {
+ Type.Repetition repetition;
+ int maxDefinitionLevel;
+ if (field.nullable()) {
+ repetition = Type.Repetition.OPTIONAL;
+ maxDefinitionLevel = 1;
+ } else {
+ repetition = Type.Repetition.REQUIRED;
+ maxDefinitionLevel = 0;
+ }
+ String[] path = new String[] {field.name()};
+
+ DataType type = field.dataType();
+
+ Types.PrimitiveBuilder builder = null;
+ // Only partition column can be `NullType`, which also uses `ConstantColumnReader`. Here we
+ // piggy-back onto Parquet boolean type for constant vector of null values, we don't really
+ // care what Parquet type it is.
+ if (type == DataTypes.BooleanType || type == DataTypes.NullType) {
+ builder = Types.primitive(PrimitiveType.PrimitiveTypeName.BOOLEAN, repetition);
+ } else if (type == DataTypes.IntegerType || type instanceof YearMonthIntervalType) {
+ builder = Types.primitive(PrimitiveType.PrimitiveTypeName.INT32, repetition);
+ } else if (type == DataTypes.DateType) {
+ builder =
+ Types.primitive(PrimitiveType.PrimitiveTypeName.INT32, repetition)
+ .as(LogicalTypeAnnotation.dateType());
+ } else if (type == DataTypes.ByteType) {
+ builder =
+ Types.primitive(PrimitiveType.PrimitiveTypeName.INT32, repetition)
+ .as(LogicalTypeAnnotation.intType(8, true));
+ } else if (type == DataTypes.ShortType) {
+ builder =
+ Types.primitive(PrimitiveType.PrimitiveTypeName.INT32, repetition)
+ .as(LogicalTypeAnnotation.intType(16, true));
+ } else if (type == DataTypes.LongType) {
+ builder = Types.primitive(PrimitiveType.PrimitiveTypeName.INT64, repetition);
+ } else if (type == DataTypes.BinaryType) {
+ builder = Types.primitive(PrimitiveType.PrimitiveTypeName.BINARY, repetition);
+ } else if (type == DataTypes.StringType) {
+ builder =
+ Types.primitive(PrimitiveType.PrimitiveTypeName.BINARY, repetition)
+ .as(LogicalTypeAnnotation.stringType());
+ } else if (type == DataTypes.FloatType) {
+ builder = Types.primitive(PrimitiveType.PrimitiveTypeName.FLOAT, repetition);
+ } else if (type == DataTypes.DoubleType) {
+ builder = Types.primitive(PrimitiveType.PrimitiveTypeName.DOUBLE, repetition);
+ } else if (type == DataTypes.TimestampType) {
+ builder =
+ Types.primitive(PrimitiveType.PrimitiveTypeName.INT64, repetition)
+ .as(LogicalTypeAnnotation.timestampType(true, TimeUnit.MICROS));
+ } else if (type == TimestampNTZType$.MODULE$) {
+ builder =
+ Types.primitive(PrimitiveType.PrimitiveTypeName.INT64, repetition)
+ .as(LogicalTypeAnnotation.timestampType(false, TimeUnit.MICROS));
+ } else if (type instanceof DecimalType) {
+ DecimalType decimalType = (DecimalType) type;
+ builder =
+ Types.primitive(PrimitiveType.PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY, repetition)
+ .length(16) // always store as Decimal128
+ .as(LogicalTypeAnnotation.decimalType(decimalType.scale(), decimalType.precision()));
+ }
+ if (builder == null) {
+ throw new UnsupportedOperationException("Unsupported input Spark type: " + type);
+ }
+
+ return new ColumnDescriptor(path, builder.named(field.name()), 0, maxDefinitionLevel);
+ }
+
+ /**
+ * Check whether the Parquet 'descriptor' and Spark read type 'sparkType' are compatible. If not,
+ * throw exception.
+ *
+ *
This mostly follows the logic in Spark's
+ * ParquetVectorUpdaterFactory#getUpdater(ColumnDescriptor, DataType)
+ *
+ * @param descriptor descriptor for a Parquet primitive column
+ * @param sparkType Spark read type
+ */
+ public static void checkParquetType(ColumnDescriptor descriptor, DataType sparkType) {
+ PrimitiveType.PrimitiveTypeName typeName = descriptor.getPrimitiveType().getPrimitiveTypeName();
+ LogicalTypeAnnotation logicalTypeAnnotation =
+ descriptor.getPrimitiveType().getLogicalTypeAnnotation();
+ boolean allowTypePromotion = (boolean) CometConf.COMET_SCHEMA_EVOLUTION_ENABLED().get();
+
+ if (sparkType instanceof NullType) {
+ return;
+ }
+
+ switch (typeName) {
+ case BOOLEAN:
+ if (sparkType == DataTypes.BooleanType) return;
+ break;
+ case INT32:
+ if (sparkType == DataTypes.IntegerType || canReadAsIntDecimal(descriptor, sparkType)) {
+ return;
+ } else if (sparkType == DataTypes.LongType
+ && isUnsignedIntTypeMatched(logicalTypeAnnotation, 32)) {
+ // In `ParquetToSparkSchemaConverter`, we map parquet UINT32 to our LongType.
+ // For unsigned int32, it stores as plain signed int32 in Parquet when dictionary
+ // fallbacks. We read them as long values.
+ return;
+ } else if (sparkType == DataTypes.LongType && allowTypePromotion) {
+ // In Comet we allow schema evolution from int to long, if
+ // `spark.comet.schemaEvolution.enabled` is enabled.
+ return;
+ } else if (sparkType == DataTypes.ByteType || sparkType == DataTypes.ShortType) {
+ return;
+ } else if (sparkType == DataTypes.DateType) {
+ // TODO: use dateTimeRebaseMode from Spark side
+ return;
+ } else if (sparkType instanceof YearMonthIntervalType) {
+ return;
+ }
+ break;
+ case INT64:
+ if (sparkType == DataTypes.LongType || canReadAsLongDecimal(descriptor, sparkType)) {
+ return;
+ } else if (isLongDecimal(sparkType)
+ && isUnsignedIntTypeMatched(logicalTypeAnnotation, 64)) {
+ // In `ParquetToSparkSchemaConverter`, we map parquet UINT64 to our Decimal(20, 0).
+ // For unsigned int64, it stores as plain signed int64 in Parquet when dictionary
+ // fallbacks. We read them as decimal values.
+ return;
+ } else if (isTimestampTypeMatched(logicalTypeAnnotation, TimeUnit.MICROS)) {
+ validateTimestampType(logicalTypeAnnotation, sparkType);
+ // TODO: use dateTimeRebaseMode from Spark side
+ return;
+ } else if (isTimestampTypeMatched(logicalTypeAnnotation, TimeUnit.MILLIS)) {
+ validateTimestampType(logicalTypeAnnotation, sparkType);
+ return;
+ }
+ break;
+ case INT96:
+ if (sparkType == TimestampNTZType$.MODULE$) {
+ convertErrorForTimestampNTZ(typeName.name());
+ } else if (sparkType == DataTypes.TimestampType) {
+ return;
+ }
+ break;
+ case FLOAT:
+ if (sparkType == DataTypes.FloatType) return;
+ // In Comet we allow schema evolution from float to double, if
+ // `spark.comet.schemaEvolution.enabled` is enabled.
+ if (sparkType == DataTypes.DoubleType && allowTypePromotion) return;
+ break;
+ case DOUBLE:
+ if (sparkType == DataTypes.DoubleType) return;
+ break;
+ case BINARY:
+ if (sparkType == DataTypes.StringType
+ || sparkType == DataTypes.BinaryType
+ || canReadAsBinaryDecimal(descriptor, sparkType)) {
+ return;
+ }
+ break;
+ case FIXED_LEN_BYTE_ARRAY:
+ if (canReadAsIntDecimal(descriptor, sparkType)
+ || canReadAsLongDecimal(descriptor, sparkType)
+ || canReadAsBinaryDecimal(descriptor, sparkType)
+ || sparkType == DataTypes.BinaryType
+ // for uuid, since iceberg maps uuid to StringType
+ || sparkType == DataTypes.StringType) {
+ return;
+ }
+ break;
+ default:
+ break;
+ }
+
+ throw new SchemaColumnConvertNotSupportedException(
+ Arrays.toString(descriptor.getPath()),
+ descriptor.getPrimitiveType().getPrimitiveTypeName().toString(),
+ sparkType.catalogString());
+ }
+
+ private static void validateTimestampType(
+ LogicalTypeAnnotation logicalTypeAnnotation, DataType sparkType) {
+ assert (logicalTypeAnnotation instanceof TimestampLogicalTypeAnnotation);
+ // Throw an exception if the Parquet type is TimestampLTZ and the Catalyst type is TimestampNTZ.
+ // This is to avoid mistakes in reading the timestamp values.
+ if (((TimestampLogicalTypeAnnotation) logicalTypeAnnotation).isAdjustedToUTC()
+ && sparkType == TimestampNTZType$.MODULE$) {
+ convertErrorForTimestampNTZ("int64 time(" + logicalTypeAnnotation + ")");
+ }
+ }
+
+ private static void convertErrorForTimestampNTZ(String parquetType) {
+ throw new RuntimeException(
+ "Unable to create Parquet converter for data type "
+ + TimestampNTZType$.MODULE$.json()
+ + " whose Parquet type is "
+ + parquetType);
+ }
+
+ private static boolean canReadAsIntDecimal(ColumnDescriptor descriptor, DataType dt) {
+ if (!DecimalType.is32BitDecimalType(dt)) return false;
+ return isDecimalTypeMatched(descriptor, dt);
+ }
+
+ private static boolean canReadAsLongDecimal(ColumnDescriptor descriptor, DataType dt) {
+ if (!DecimalType.is64BitDecimalType(dt)) return false;
+ return isDecimalTypeMatched(descriptor, dt);
+ }
+
+ private static boolean canReadAsBinaryDecimal(ColumnDescriptor descriptor, DataType dt) {
+ if (!DecimalType.isByteArrayDecimalType(dt)) return false;
+ return isDecimalTypeMatched(descriptor, dt);
+ }
+
+ private static boolean isLongDecimal(DataType dt) {
+ if (dt instanceof DecimalType) {
+ DecimalType d = (DecimalType) dt;
+ return d.precision() == 20 && d.scale() == 0;
+ }
+ return false;
+ }
+
+ private static boolean isDecimalTypeMatched(ColumnDescriptor descriptor, DataType dt) {
+ DecimalType d = (DecimalType) dt;
+ LogicalTypeAnnotation typeAnnotation = descriptor.getPrimitiveType().getLogicalTypeAnnotation();
+ if (typeAnnotation instanceof DecimalLogicalTypeAnnotation) {
+ DecimalLogicalTypeAnnotation decimalType = (DecimalLogicalTypeAnnotation) typeAnnotation;
+ // It's OK if the required decimal precision is larger than or equal to the physical decimal
+ // precision in the Parquet metadata, as long as the decimal scale is the same.
+ return decimalType.getPrecision() <= d.precision() && decimalType.getScale() == d.scale();
+ }
+ return false;
+ }
+
+ private static boolean isTimestampTypeMatched(
+ LogicalTypeAnnotation logicalTypeAnnotation, LogicalTypeAnnotation.TimeUnit unit) {
+ return logicalTypeAnnotation instanceof TimestampLogicalTypeAnnotation
+ && ((TimestampLogicalTypeAnnotation) logicalTypeAnnotation).getUnit() == unit;
+ }
+
+ private static boolean isUnsignedIntTypeMatched(
+ LogicalTypeAnnotation logicalTypeAnnotation, int bitWidth) {
+ return logicalTypeAnnotation instanceof IntLogicalTypeAnnotation
+ && !((IntLogicalTypeAnnotation) logicalTypeAnnotation).isSigned()
+ && ((IntLogicalTypeAnnotation) logicalTypeAnnotation).getBitWidth() == bitWidth;
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/parquet/Utils.java b/common/src/main/java/org/apache/comet/parquet/Utils.java
new file mode 100644
index 000000000..95ca06cda
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/parquet/Utils.java
@@ -0,0 +1,238 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.schema.LogicalTypeAnnotation;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.spark.sql.types.*;
+
+public class Utils {
+ public static ColumnReader getColumnReader(
+ DataType type,
+ ColumnDescriptor descriptor,
+ int batchSize,
+ boolean useDecimal128,
+ boolean useLazyMaterialization) {
+ // TODO: support `useLegacyDateTimestamp` for Iceberg
+ return getColumnReader(
+ type, descriptor, batchSize, useDecimal128, useLazyMaterialization, true);
+ }
+
+ public static ColumnReader getColumnReader(
+ DataType type,
+ ColumnDescriptor descriptor,
+ int batchSize,
+ boolean useDecimal128,
+ boolean useLazyMaterialization,
+ boolean useLegacyDateTimestamp) {
+ if (useLazyMaterialization && supportLazyMaterialization(type)) {
+ return new LazyColumnReader(
+ type, descriptor, batchSize, useDecimal128, useLegacyDateTimestamp);
+ } else {
+ return new ColumnReader(type, descriptor, batchSize, useDecimal128, useLegacyDateTimestamp);
+ }
+ }
+
+ private static boolean supportLazyMaterialization(DataType type) {
+ return (type instanceof StringType || type instanceof BinaryType);
+ }
+
+ /**
+ * Initialize the Comet native Parquet reader.
+ *
+ * @param descriptor the Parquet column descriptor for the column to be read
+ * @param readType the Spark read type used for type promotion. Null if promotion is not enabled.
+ * @param batchSize the batch size, i.e., maximum number of elements per record batch
+ * @param useDecimal128 whether to always represent decimals using 128 bits. If false, the native
+ * reader may represent decimals using 32 or 64 bits, depending on the precision.
+ * @param useLegacyDateTimestampOrNTZ whether to read dates/timestamps that were written in the
+ * legacy hybrid Julian + Gregorian calendar as it is. If false, throw exceptions instead. If
+ * the spark type is TimestampNTZ, this should be true.
+ */
+ public static long initColumnReader(
+ ColumnDescriptor descriptor,
+ DataType readType,
+ int batchSize,
+ boolean useDecimal128,
+ boolean useLegacyDateTimestampOrNTZ) {
+ PrimitiveType primitiveType = descriptor.getPrimitiveType();
+ int primitiveTypeId = getPhysicalTypeId(primitiveType.getPrimitiveTypeName());
+ LogicalTypeAnnotation annotation = primitiveType.getLogicalTypeAnnotation();
+
+ // Process logical type information
+
+ int bitWidth = -1;
+ boolean isSigned = false;
+ if (annotation instanceof LogicalTypeAnnotation.IntLogicalTypeAnnotation) {
+ LogicalTypeAnnotation.IntLogicalTypeAnnotation intAnnotation =
+ (LogicalTypeAnnotation.IntLogicalTypeAnnotation) annotation;
+ bitWidth = intAnnotation.getBitWidth();
+ isSigned = intAnnotation.isSigned();
+ }
+
+ int precision, scale;
+ precision = scale = -1;
+ if (annotation instanceof LogicalTypeAnnotation.DecimalLogicalTypeAnnotation) {
+ LogicalTypeAnnotation.DecimalLogicalTypeAnnotation decimalAnnotation =
+ (LogicalTypeAnnotation.DecimalLogicalTypeAnnotation) annotation;
+ precision = decimalAnnotation.getPrecision();
+ scale = decimalAnnotation.getScale();
+ }
+
+ int tu = -1;
+ boolean isAdjustedUtc = false;
+ if (annotation instanceof LogicalTypeAnnotation.TimestampLogicalTypeAnnotation) {
+ LogicalTypeAnnotation.TimestampLogicalTypeAnnotation timestampAnnotation =
+ (LogicalTypeAnnotation.TimestampLogicalTypeAnnotation) annotation;
+ tu = getTimeUnitId(timestampAnnotation.getUnit());
+ isAdjustedUtc = timestampAnnotation.isAdjustedToUTC();
+ }
+
+ TypePromotionInfo promotionInfo;
+ if (readType != null) {
+ promotionInfo = new TypePromotionInfo(readType);
+ } else {
+ // If type promotion is not enable, we'll just use the Parquet primitive type and precision.
+ promotionInfo = new TypePromotionInfo(primitiveTypeId, precision);
+ }
+
+ return Native.initColumnReader(
+ primitiveTypeId,
+ getLogicalTypeId(annotation),
+ promotionInfo.physicalTypeId,
+ descriptor.getPath(),
+ descriptor.getMaxDefinitionLevel(),
+ descriptor.getMaxRepetitionLevel(),
+ bitWidth,
+ isSigned,
+ primitiveType.getTypeLength(),
+ precision,
+ promotionInfo.precision,
+ scale,
+ tu,
+ isAdjustedUtc,
+ batchSize,
+ useDecimal128,
+ useLegacyDateTimestampOrNTZ);
+ }
+
+ static class TypePromotionInfo {
+ // The Parquet physical type ID converted from the Spark read schema, or the original Parquet
+ // physical type ID if type promotion is not enabled.
+ int physicalTypeId;
+ // Decimal precision from the Spark read schema, or -1 if it's not decimal type.
+ int precision;
+
+ TypePromotionInfo(int physicalTypeId, int precision) {
+ this.physicalTypeId = physicalTypeId;
+ this.precision = precision;
+ }
+
+ TypePromotionInfo(DataType sparkReadType) {
+ // Create a dummy `StructField` from the input Spark type. We don't care about
+ // field name, nullability and metadata.
+ StructField f = new StructField("f", sparkReadType, false, Metadata.empty());
+ ColumnDescriptor descriptor = TypeUtil.convertToParquet(f);
+ PrimitiveType primitiveType = descriptor.getPrimitiveType();
+ int physicalTypeId = getPhysicalTypeId(primitiveType.getPrimitiveTypeName());
+ LogicalTypeAnnotation annotation = primitiveType.getLogicalTypeAnnotation();
+ int precision = -1;
+ if (annotation instanceof LogicalTypeAnnotation.DecimalLogicalTypeAnnotation) {
+ LogicalTypeAnnotation.DecimalLogicalTypeAnnotation decimalAnnotation =
+ (LogicalTypeAnnotation.DecimalLogicalTypeAnnotation) annotation;
+ precision = decimalAnnotation.getPrecision();
+ }
+ this.physicalTypeId = physicalTypeId;
+ this.precision = precision;
+ }
+ }
+
+ /**
+ * Maps the input Parquet physical type 'typeName' to an integer representing it. This is used for
+ * serialization between the Java and native side.
+ *
+ * @param typeName enum for the Parquet physical type
+ * @return an integer representing the input physical type
+ */
+ static int getPhysicalTypeId(PrimitiveType.PrimitiveTypeName typeName) {
+ switch (typeName) {
+ case BOOLEAN:
+ return 0;
+ case INT32:
+ return 1;
+ case INT64:
+ return 2;
+ case INT96:
+ return 3;
+ case FLOAT:
+ return 4;
+ case DOUBLE:
+ return 5;
+ case BINARY:
+ return 6;
+ case FIXED_LEN_BYTE_ARRAY:
+ return 7;
+ }
+ throw new IllegalArgumentException("Invalid Parquet physical type: " + typeName);
+ }
+
+ /**
+ * Maps the input Parquet logical type 'annotation' to an integer representing it. This is used
+ * for serialization between the Java and native side.
+ *
+ * @param annotation the Parquet logical type annotation
+ * @return an integer representing the input logical type
+ */
+ static int getLogicalTypeId(LogicalTypeAnnotation annotation) {
+ if (annotation == null) {
+ return -1; // No logical type associated
+ } else if (annotation instanceof LogicalTypeAnnotation.IntLogicalTypeAnnotation) {
+ return 0;
+ } else if (annotation instanceof LogicalTypeAnnotation.StringLogicalTypeAnnotation) {
+ return 1;
+ } else if (annotation instanceof LogicalTypeAnnotation.DecimalLogicalTypeAnnotation) {
+ return 2;
+ } else if (annotation instanceof LogicalTypeAnnotation.DateLogicalTypeAnnotation) {
+ return 3;
+ } else if (annotation instanceof LogicalTypeAnnotation.TimestampLogicalTypeAnnotation) {
+ return 4;
+ } else if (annotation instanceof LogicalTypeAnnotation.EnumLogicalTypeAnnotation) {
+ return 5;
+ } else if (annotation instanceof LogicalTypeAnnotation.UUIDLogicalTypeAnnotation) {
+ return 6;
+ }
+
+ throw new UnsupportedOperationException("Unsupported Parquet logical type " + annotation);
+ }
+
+ static int getTimeUnitId(LogicalTypeAnnotation.TimeUnit tu) {
+ switch (tu) {
+ case MILLIS:
+ return 0;
+ case MICROS:
+ return 1;
+ case NANOS:
+ return 2;
+ default:
+ throw new UnsupportedOperationException("Unsupported TimeUnit " + tu);
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometDecodedVector.java b/common/src/main/java/org/apache/comet/vector/CometDecodedVector.java
new file mode 100644
index 000000000..fec266bf1
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometDecodedVector.java
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import org.apache.arrow.vector.BaseVariableWidthVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.spark.sql.comet.util.Utils;
+import org.apache.spark.unsafe.Platform;
+
+/** A Comet vector whose elements are already decoded (i.e., materialized). */
+public abstract class CometDecodedVector extends CometVector {
+ /**
+ * The vector that stores all the values. For dictionary-backed vector, this is the vector of
+ * indices.
+ */
+ protected final ValueVector valueVector;
+
+ private boolean hasNull;
+ private int numNulls;
+ private int numValues;
+ private int validityByteCacheIndex = -1;
+ private byte validityByteCache;
+
+ protected CometDecodedVector(ValueVector vector, Field valueField, boolean useDecimal128) {
+ super(Utils.fromArrowField(valueField), useDecimal128);
+ this.valueVector = vector;
+ this.numNulls = valueVector.getNullCount();
+ this.numValues = valueVector.getValueCount();
+ this.hasNull = numNulls != 0;
+ }
+
+ @Override
+ ValueVector getValueVector() {
+ return valueVector;
+ }
+
+ @Override
+ public void setNumNulls(int numNulls) {
+ // We don't need to update null count in 'valueVector' since 'ValueVector.getNullCount' will
+ // re-compute the null count from validity buffer.
+ this.numNulls = numNulls;
+ this.hasNull = numNulls != 0;
+ this.validityByteCacheIndex = -1;
+ }
+
+ @Override
+ public void setNumValues(int numValues) {
+ this.numValues = numValues;
+ if (valueVector instanceof BaseVariableWidthVector) {
+ BaseVariableWidthVector bv = (BaseVariableWidthVector) valueVector;
+ // In case `lastSet` is smaller than `numValues`, `setValueCount` will set all the offsets
+ // within `[lastSet + 1, numValues)` to be empty, which is incorrect in our case.
+ //
+ // For instance, this can happen if one first call `setNumValues` with input 100, and then
+ // again `setNumValues` with 200. The first call will set `lastSet` to 99, while the second
+ // call will set all strings between indices `[100, 200)` to be empty.
+ bv.setLastSet(numValues);
+ }
+ valueVector.setValueCount(numValues);
+ }
+
+ public int numValues() {
+ return numValues;
+ }
+
+ @Override
+ public boolean hasNull() {
+ return hasNull;
+ }
+
+ @Override
+ public int numNulls() {
+ return numNulls;
+ }
+
+ @Override
+ public boolean isNullAt(int rowId) {
+ if (!hasNull) return false;
+
+ int byteIndex = rowId >> 3;
+ if (byteIndex != validityByteCacheIndex) {
+ long validityBufferAddress = valueVector.getValidityBuffer().memoryAddress();
+ validityByteCache = Platform.getByte(null, validityBufferAddress + byteIndex);
+ validityByteCacheIndex = byteIndex;
+ }
+ return ((validityByteCache >> (rowId & 7)) & 1) == 0;
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometDelegateVector.java b/common/src/main/java/org/apache/comet/vector/CometDelegateVector.java
new file mode 100644
index 000000000..165af5631
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometDelegateVector.java
@@ -0,0 +1,169 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarArray;
+import org.apache.spark.sql.vectorized.ColumnarMap;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/** A special Comet vector that just delegate all calls */
+public class CometDelegateVector extends CometVector {
+ protected CometVector delegate;
+
+ public CometDelegateVector(DataType dataType) {
+ this(dataType, null, false);
+ }
+
+ public CometDelegateVector(DataType dataType, boolean useDecimal128) {
+ this(dataType, null, useDecimal128);
+ }
+
+ public CometDelegateVector(DataType dataType, CometVector delegate, boolean useDecimal128) {
+ super(dataType, useDecimal128);
+ if (delegate instanceof CometDelegateVector) {
+ throw new IllegalArgumentException("cannot have nested delegation");
+ }
+ this.delegate = delegate;
+ }
+
+ public void setDelegate(CometVector delegate) {
+ this.delegate = delegate;
+ }
+
+ @Override
+ public void setNumNulls(int numNulls) {
+ delegate.setNumNulls(numNulls);
+ }
+
+ @Override
+ public void setNumValues(int numValues) {
+ delegate.setNumValues(numValues);
+ }
+
+ @Override
+ public int numValues() {
+ return delegate.numValues();
+ }
+
+ @Override
+ public boolean hasNull() {
+ return delegate.hasNull();
+ }
+
+ @Override
+ public int numNulls() {
+ return delegate.numNulls();
+ }
+
+ @Override
+ public boolean isNullAt(int rowId) {
+ return delegate.isNullAt(rowId);
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ return delegate.getBoolean(rowId);
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return delegate.getByte(rowId);
+ }
+
+ @Override
+ public short getShort(int rowId) {
+ return delegate.getShort(rowId);
+ }
+
+ @Override
+ public int getInt(int rowId) {
+ return delegate.getInt(rowId);
+ }
+
+ @Override
+ public long getLong(int rowId) {
+ return delegate.getLong(rowId);
+ }
+
+ @Override
+ public float getFloat(int rowId) {
+ return delegate.getFloat(rowId);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return delegate.getDouble(rowId);
+ }
+
+ @Override
+ public Decimal getDecimal(int i, int precision, int scale) {
+ return delegate.getDecimal(i, precision, scale);
+ }
+
+ @Override
+ byte[] getBinaryDecimal(int i) {
+ return delegate.getBinaryDecimal(i);
+ }
+
+ @Override
+ public UTF8String getUTF8String(int rowId) {
+ return delegate.getUTF8String(rowId);
+ }
+
+ @Override
+ public byte[] getBinary(int rowId) {
+ return delegate.getBinary(rowId);
+ }
+
+ @Override
+ public ColumnarArray getArray(int i) {
+ return delegate.getArray(i);
+ }
+
+ @Override
+ public ColumnarMap getMap(int i) {
+ return delegate.getMap(i);
+ }
+
+ @Override
+ public ColumnVector getChild(int i) {
+ return delegate.getChild(i);
+ }
+
+ @Override
+ ValueVector getValueVector() {
+ return delegate.getValueVector();
+ }
+
+ @Override
+ public CometVector slice(int offset, int length) {
+ return delegate.slice(offset, length);
+ }
+
+ @Override
+ DictionaryProvider getDictionaryProvider() {
+ return delegate.getDictionaryProvider();
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometDictionary.java b/common/src/main/java/org/apache/comet/vector/CometDictionary.java
new file mode 100644
index 000000000..8c7046585
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometDictionary.java
@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import org.apache.arrow.vector.ValueVector;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/** A dictionary which maps indices (integers) to values. */
+public class CometDictionary implements AutoCloseable {
+ private static final int DECIMAL_BYTE_WIDTH = 16;
+
+ private final CometPlainVector values;
+ private final int numValues;
+
+ /** Decoded dictionary values. Only one of the following is set. */
+ private byte[] bytes;
+
+ private short[] shorts;
+ private int[] ints;
+ private long[] longs;
+ private float[] floats;
+ private double[] doubles;
+ private boolean[] booleans;
+ private ByteArrayWrapper[] binaries;
+ private UTF8String[] strings;
+
+ public CometDictionary(CometPlainVector values) {
+ this.values = values;
+ this.numValues = values.numValues();
+ initialize();
+ }
+
+ public ValueVector getValueVector() {
+ return values.getValueVector();
+ }
+
+ public boolean decodeToBoolean(int index) {
+ return booleans[index];
+ }
+
+ public byte decodeToByte(int index) {
+ return bytes[index];
+ }
+
+ public short decodeToShort(int index) {
+ return shorts[index];
+ }
+
+ public int decodeToInt(int index) {
+ return ints[index];
+ }
+
+ public long decodeToLong(int index) {
+ return longs[index];
+ }
+
+ public float decodeToFloat(int index) {
+ return floats[index];
+ }
+
+ public double decodeToDouble(int index) {
+ return doubles[index];
+ }
+
+ public byte[] decodeToBinary(int index) {
+ return binaries[index].bytes;
+ }
+
+ public UTF8String decodeToUTF8String(int index) {
+ return strings[index];
+ }
+
+ @Override
+ public void close() {
+ values.close();
+ }
+
+ private void initialize() {
+ switch (values.getValueVector().getMinorType()) {
+ case BIT:
+ booleans = new boolean[numValues];
+ for (int i = 0; i < numValues; i++) {
+ booleans[i] = values.getBoolean(i);
+ }
+ break;
+ case TINYINT:
+ bytes = new byte[numValues];
+ for (int i = 0; i < numValues; i++) {
+ bytes[i] = values.getByte(i);
+ }
+ break;
+ case SMALLINT:
+ shorts = new short[numValues];
+ for (int i = 0; i < numValues; i++) {
+ shorts[i] = values.getShort(i);
+ }
+ break;
+ case INT:
+ case DATEDAY:
+ ints = new int[numValues];
+ for (int i = 0; i < numValues; i++) {
+ ints[i] = values.getInt(i);
+ }
+ break;
+ case BIGINT:
+ case TIMESTAMPMICRO:
+ case TIMESTAMPMICROTZ:
+ longs = new long[numValues];
+ for (int i = 0; i < numValues; i++) {
+ longs[i] = values.getLong(i);
+ }
+ break;
+ case FLOAT4:
+ floats = new float[numValues];
+ for (int i = 0; i < numValues; i++) {
+ floats[i] = values.getFloat(i);
+ }
+ break;
+ case FLOAT8:
+ doubles = new double[numValues];
+ for (int i = 0; i < numValues; i++) {
+ doubles[i] = values.getDouble(i);
+ }
+ break;
+ case VARBINARY:
+ case FIXEDSIZEBINARY:
+ binaries = new ByteArrayWrapper[numValues];
+ for (int i = 0; i < numValues; i++) {
+ binaries[i] = new ByteArrayWrapper(values.getBinary(i));
+ }
+ break;
+ case VARCHAR:
+ strings = new UTF8String[numValues];
+ for (int i = 0; i < numValues; i++) {
+ strings[i] = values.getUTF8String(i);
+ }
+ break;
+ case DECIMAL:
+ binaries = new ByteArrayWrapper[numValues];
+ for (int i = 0; i < numValues; i++) {
+ // Need copying here since we re-use byte array for decimal
+ byte[] bytes = values.getBinaryDecimal(i);
+ byte[] copy = new byte[DECIMAL_BYTE_WIDTH];
+ System.arraycopy(bytes, 0, copy, 0, DECIMAL_BYTE_WIDTH);
+ binaries[i] = new ByteArrayWrapper(copy);
+ }
+ break;
+ default:
+ throw new IllegalArgumentException(
+ "Invalid Arrow minor type: " + values.getValueVector().getMinorType());
+ }
+ }
+
+ private static class ByteArrayWrapper {
+ private final byte[] bytes;
+
+ ByteArrayWrapper(byte[] bytes) {
+ this.bytes = bytes;
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometDictionaryVector.java b/common/src/main/java/org/apache/comet/vector/CometDictionaryVector.java
new file mode 100644
index 000000000..225fcfc43
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometDictionaryVector.java
@@ -0,0 +1,135 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.util.TransferPair;
+import org.apache.parquet.Preconditions;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/** A column vector whose elements are dictionary-encoded. */
+public class CometDictionaryVector extends CometDecodedVector {
+ public final CometPlainVector indices;
+ public final CometDictionary values;
+ public final DictionaryProvider provider;
+
+ /** Whether this vector is an alias sliced from another vector. */
+ private final boolean isAlias;
+
+ public CometDictionaryVector(
+ CometPlainVector indices,
+ CometDictionary values,
+ DictionaryProvider provider,
+ boolean useDecimal128) {
+ this(indices, values, provider, useDecimal128, false);
+ }
+
+ public CometDictionaryVector(
+ CometPlainVector indices,
+ CometDictionary values,
+ DictionaryProvider provider,
+ boolean useDecimal128,
+ boolean isAlias) {
+ super(indices.valueVector, values.getValueVector().getField(), useDecimal128);
+ Preconditions.checkArgument(
+ indices.valueVector instanceof IntVector, "'indices' should be a IntVector");
+ this.values = values;
+ this.indices = indices;
+ this.provider = provider;
+ this.isAlias = isAlias;
+ }
+
+ @Override
+ DictionaryProvider getDictionaryProvider() {
+ return this.provider;
+ }
+
+ @Override
+ public void close() {
+ super.close();
+ // Only close the values vector if this is not a sliced vector.
+ if (!isAlias) {
+ values.close();
+ }
+ }
+
+ @Override
+ public boolean getBoolean(int i) {
+ return values.decodeToBoolean(indices.getInt(i));
+ }
+
+ @Override
+ public byte getByte(int i) {
+ return values.decodeToByte(indices.getInt(i));
+ }
+
+ @Override
+ public short getShort(int i) {
+ return values.decodeToShort(indices.getInt(i));
+ }
+
+ @Override
+ public int getInt(int i) {
+ return values.decodeToInt(indices.getInt(i));
+ }
+
+ @Override
+ public long getLong(int i) {
+ return values.decodeToLong(indices.getInt(i));
+ }
+
+ @Override
+ public float getFloat(int i) {
+ return values.decodeToFloat(indices.getInt(i));
+ }
+
+ @Override
+ public double getDouble(int i) {
+ return values.decodeToDouble(indices.getInt(i));
+ }
+
+ @Override
+ public UTF8String getUTF8String(int i) {
+ return values.decodeToUTF8String(indices.getInt(i));
+ }
+
+ @Override
+ public byte[] getBinary(int i) {
+ return values.decodeToBinary(indices.getInt(i));
+ }
+
+ @Override
+ byte[] getBinaryDecimal(int i) {
+ return values.decodeToBinary(indices.getInt(i));
+ }
+
+ @Override
+ public CometVector slice(int offset, int length) {
+ TransferPair tp = indices.valueVector.getTransferPair(indices.valueVector.getAllocator());
+ tp.splitAndTransfer(offset, length);
+ CometPlainVector sliced = new CometPlainVector(tp.getTo(), useDecimal128);
+
+ // Set the alias flag to true so that the sliced vector will not close the dictionary vector.
+ // Otherwise, if the dictionary is closed, the sliced vector will not be able to access the
+ // dictionary.
+ return new CometDictionaryVector(sliced, values, provider, useDecimal128, true);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometLazyVector.java b/common/src/main/java/org/apache/comet/vector/CometLazyVector.java
new file mode 100644
index 000000000..17b8d7e71
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometLazyVector.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import org.apache.arrow.vector.ValueVector;
+import org.apache.spark.sql.types.DataType;
+
+import org.apache.comet.parquet.LazyColumnReader;
+
+public class CometLazyVector extends CometDelegateVector {
+ private final LazyColumnReader columnReader;
+
+ public CometLazyVector(DataType type, LazyColumnReader columnReader, boolean useDecimal128) {
+ super(type, useDecimal128);
+ this.columnReader = columnReader;
+ }
+
+ public CometDecodedVector getDecodedVector() {
+ return (CometDecodedVector) delegate;
+ }
+
+ @Override
+ public ValueVector getValueVector() {
+ columnReader.readAllBatch();
+ setDelegate(columnReader.loadVector());
+ return super.getValueVector();
+ }
+
+ @Override
+ public void setNumNulls(int numNulls) {
+ throw new UnsupportedOperationException("CometLazyVector doesn't support 'setNumNulls'");
+ }
+
+ @Override
+ public void setNumValues(int numValues) {
+ throw new UnsupportedOperationException("CometLazyVector doesn't support 'setNumValues'");
+ }
+
+ @Override
+ public void close() {
+ // Do nothing. 'vector' is closed by 'columnReader' which owns it.
+ }
+
+ @Override
+ public boolean hasNull() {
+ columnReader.readAllBatch();
+ setDelegate(columnReader.loadVector());
+ return super.hasNull();
+ }
+
+ @Override
+ public int numNulls() {
+ columnReader.readAllBatch();
+ setDelegate(columnReader.loadVector());
+ return super.numNulls();
+ }
+
+ @Override
+ public boolean isNullAt(int rowId) {
+ if (columnReader.materializeUpToIfNecessary(rowId)) {
+ setDelegate(columnReader.loadVector());
+ }
+ return super.isNullAt(rowId);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometListVector.java b/common/src/main/java/org/apache/comet/vector/CometListVector.java
new file mode 100644
index 000000000..1c8f3e658
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometListVector.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import org.apache.arrow.vector.*;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.util.TransferPair;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarArray;
+
+/** A Comet column vector for list type. */
+public class CometListVector extends CometDecodedVector {
+ final ListVector listVector;
+ final ValueVector dataVector;
+ final ColumnVector dataColumnVector;
+
+ public CometListVector(ValueVector vector, boolean useDecimal128) {
+ super(vector, vector.getField(), useDecimal128);
+
+ this.listVector = ((ListVector) vector);
+ this.dataVector = listVector.getDataVector();
+ this.dataColumnVector = getVector(dataVector, useDecimal128);
+ }
+
+ @Override
+ public ColumnarArray getArray(int i) {
+ int start = listVector.getOffsetBuffer().getInt(i * ListVector.OFFSET_WIDTH);
+ int end = listVector.getOffsetBuffer().getInt((i + 1) * ListVector.OFFSET_WIDTH);
+
+ return new ColumnarArray(dataColumnVector, start, end - start);
+ }
+
+ @Override
+ public CometVector slice(int offset, int length) {
+ TransferPair tp = this.valueVector.getTransferPair(this.valueVector.getAllocator());
+ tp.splitAndTransfer(offset, length);
+
+ return new CometListVector(tp.getTo(), useDecimal128);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometPlainVector.java b/common/src/main/java/org/apache/comet/vector/CometPlainVector.java
new file mode 100644
index 000000000..a7373590a
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometPlainVector.java
@@ -0,0 +1,172 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.UUID;
+
+import org.apache.arrow.c.CDataDictionaryProvider;
+import org.apache.arrow.vector.*;
+import org.apache.arrow.vector.util.TransferPair;
+import org.apache.parquet.Preconditions;
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/** A column vector whose elements are plainly decoded. */
+public class CometPlainVector extends CometDecodedVector {
+ private final long valueBufferAddress;
+ private final boolean isBaseFixedWidthVector;
+
+ private byte booleanByteCache;
+ private int booleanByteCacheIndex = -1;
+
+ public CometPlainVector(ValueVector vector, boolean useDecimal128) {
+ super(vector, vector.getField(), useDecimal128);
+ // NullType doesn't have data buffer.
+ if (vector instanceof NullVector) {
+ this.valueBufferAddress = -1;
+ } else {
+ this.valueBufferAddress = vector.getDataBuffer().memoryAddress();
+ }
+
+ isBaseFixedWidthVector = valueVector instanceof BaseFixedWidthVector;
+ }
+
+ @Override
+ public void setNumNulls(int numNulls) {
+ super.setNumNulls(numNulls);
+ this.booleanByteCacheIndex = -1;
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ int byteIndex = rowId >> 3;
+ if (byteIndex != booleanByteCacheIndex) {
+ booleanByteCache = getByte(byteIndex);
+ booleanByteCacheIndex = byteIndex;
+ }
+ return ((booleanByteCache >> (rowId & 7)) & 1) == 1;
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ return Platform.getByte(null, valueBufferAddress + rowId);
+ }
+
+ @Override
+ public short getShort(int rowId) {
+ return Platform.getShort(null, valueBufferAddress + rowId * 2L);
+ }
+
+ @Override
+ public int getInt(int rowId) {
+ return Platform.getInt(null, valueBufferAddress + rowId * 4L);
+ }
+
+ @Override
+ public long getLong(int rowId) {
+ return Platform.getLong(null, valueBufferAddress + rowId * 8L);
+ }
+
+ @Override
+ public float getFloat(int rowId) {
+ return Platform.getFloat(null, valueBufferAddress + rowId * 4L);
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ return Platform.getDouble(null, valueBufferAddress + rowId * 8L);
+ }
+
+ @Override
+ public UTF8String getUTF8String(int rowId) {
+ if (!isBaseFixedWidthVector) {
+ BaseVariableWidthVector varWidthVector = (BaseVariableWidthVector) valueVector;
+ long offsetBufferAddress = varWidthVector.getOffsetBuffer().memoryAddress();
+ int offset = Platform.getInt(null, offsetBufferAddress + rowId * 4L);
+ int length = Platform.getInt(null, offsetBufferAddress + (rowId + 1L) * 4L) - offset;
+ return UTF8String.fromAddress(null, valueBufferAddress + offset, length);
+ } else {
+ // Iceberg maps UUID to StringType.
+ // The data type here must be UUID because the only FLBA -> String mapping we have is UUID.
+ BaseFixedWidthVector fixedWidthVector = (BaseFixedWidthVector) valueVector;
+ int length = fixedWidthVector.getTypeWidth();
+ int offset = rowId * length;
+ byte[] result = new byte[length];
+ Platform.copyMemory(
+ null, valueBufferAddress + offset, result, Platform.BYTE_ARRAY_OFFSET, length);
+ return UTF8String.fromString(convertToUuid(result).toString());
+ }
+ }
+
+ @Override
+ public byte[] getBinary(int rowId) {
+ int offset;
+ int length;
+ if (valueVector instanceof BaseVariableWidthVector) {
+ BaseVariableWidthVector varWidthVector = (BaseVariableWidthVector) valueVector;
+ long offsetBufferAddress = varWidthVector.getOffsetBuffer().memoryAddress();
+ offset = Platform.getInt(null, offsetBufferAddress + rowId * 4L);
+ length = Platform.getInt(null, offsetBufferAddress + (rowId + 1L) * 4L) - offset;
+ } else if (valueVector instanceof BaseFixedWidthVector) {
+ BaseFixedWidthVector fixedWidthVector = (BaseFixedWidthVector) valueVector;
+ length = fixedWidthVector.getTypeWidth();
+ offset = rowId * length;
+ } else {
+ throw new RuntimeException("Unsupported binary vector type: " + valueVector.getName());
+ }
+ byte[] result = new byte[length];
+ Platform.copyMemory(
+ null, valueBufferAddress + offset, result, Platform.BYTE_ARRAY_OFFSET, length);
+ return result;
+ }
+
+ @Override
+ CDataDictionaryProvider getDictionaryProvider() {
+ return null;
+ }
+
+ @Override
+ public boolean isNullAt(int rowId) {
+ if (this.valueBufferAddress == -1) {
+ return true;
+ } else {
+ return super.isNullAt(rowId);
+ }
+ }
+
+ @Override
+ public CometVector slice(int offset, int length) {
+ TransferPair tp = this.valueVector.getTransferPair(this.valueVector.getAllocator());
+ tp.splitAndTransfer(offset, length);
+
+ return new CometPlainVector(tp.getTo(), useDecimal128);
+ }
+
+ private static UUID convertToUuid(byte[] buf) {
+ Preconditions.checkArgument(buf.length == 16, "UUID require 16 bytes");
+ ByteBuffer bb = ByteBuffer.wrap(buf);
+ bb.order(ByteOrder.BIG_ENDIAN);
+ long mostSigBits = bb.getLong();
+ long leastSigBits = bb.getLong();
+ return new UUID(mostSigBits, leastSigBits);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometStructVector.java b/common/src/main/java/org/apache/comet/vector/CometStructVector.java
new file mode 100644
index 000000000..52dcd4326
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometStructVector.java
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.arrow.vector.*;
+import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.util.TransferPair;
+import org.apache.spark.sql.vectorized.ColumnVector;
+
+/** A Comet column vector for struct type. */
+public class CometStructVector extends CometDecodedVector {
+ final List children;
+
+ public CometStructVector(ValueVector vector, boolean useDecimal128) {
+ super(vector, vector.getField(), useDecimal128);
+
+ StructVector structVector = ((StructVector) vector);
+
+ int size = structVector.size();
+ List children = new ArrayList<>();
+
+ for (int i = 0; i < size; ++i) {
+ ValueVector value = structVector.getVectorById(i);
+ children.add(getVector(value, useDecimal128));
+ }
+ this.children = children;
+ }
+
+ @Override
+ public ColumnVector getChild(int i) {
+ return children.get(i);
+ }
+
+ @Override
+ public CometVector slice(int offset, int length) {
+ TransferPair tp = this.valueVector.getTransferPair(this.valueVector.getAllocator());
+ tp.splitAndTransfer(offset, length);
+
+ return new CometStructVector(tp.getTo(), useDecimal128);
+ }
+}
diff --git a/common/src/main/java/org/apache/comet/vector/CometVector.java b/common/src/main/java/org/apache/comet/vector/CometVector.java
new file mode 100644
index 000000000..744f24395
--- /dev/null
+++ b/common/src/main/java/org/apache/comet/vector/CometVector.java
@@ -0,0 +1,232 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector;
+
+import java.math.BigDecimal;
+import java.math.BigInteger;
+
+import org.apache.arrow.vector.FixedWidthVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.complex.ListVector;
+import org.apache.arrow.vector.complex.StructVector;
+import org.apache.arrow.vector.dictionary.Dictionary;
+import org.apache.arrow.vector.dictionary.DictionaryProvider;
+import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.types.IntegerType;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarArray;
+import org.apache.spark.sql.vectorized.ColumnarMap;
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/** Base class for all Comet column vector implementations. */
+public abstract class CometVector extends ColumnVector {
+ private static final int DECIMAL_BYTE_WIDTH = 16;
+ private final byte[] DECIMAL_BYTES = new byte[DECIMAL_BYTE_WIDTH];
+ protected final boolean useDecimal128;
+
+ protected CometVector(DataType type, boolean useDecimal128) {
+ super(type);
+ this.useDecimal128 = useDecimal128;
+ }
+
+ /**
+ * Sets the number of nulls in this vector to be 'numNulls'. This is used when the vector is
+ * reused across batches.
+ */
+ public abstract void setNumNulls(int numNulls);
+
+ /**
+ * Sets the number of values (including both nulls and non-nulls) in this vector to be
+ * 'numValues'. This is used when the vector is reused across batches.
+ */
+ public abstract void setNumValues(int numValues);
+
+ /** Returns the number of values in this vector. */
+ public abstract int numValues();
+
+ /** Whether the elements of this vector are of fixed length. */
+ public boolean isFixedLength() {
+ return getValueVector() instanceof FixedWidthVector;
+ }
+
+ @Override
+ public Decimal getDecimal(int i, int precision, int scale) {
+ if (!useDecimal128 && precision <= Decimal.MAX_INT_DIGITS() && type instanceof IntegerType) {
+ return Decimal.createUnsafe(getInt(i), precision, scale);
+ } else if (!useDecimal128 && precision <= Decimal.MAX_LONG_DIGITS()) {
+ return Decimal.createUnsafe(getLong(i), precision, scale);
+ } else {
+ byte[] bytes = getBinaryDecimal(i);
+ BigInteger bigInteger = new BigInteger(bytes);
+ BigDecimal javaDecimal = new BigDecimal(bigInteger, scale);
+ try {
+ return Decimal.apply(javaDecimal, precision, scale);
+ } catch (ArithmeticException e) {
+ throw new ArithmeticException(
+ "Cannot convert "
+ + javaDecimal
+ + " (bytes: "
+ + bytes
+ + ", integer: "
+ + bigInteger
+ + ") to decimal with precision: "
+ + precision
+ + " and scale: "
+ + scale);
+ }
+ }
+ }
+
+ /** Reads a 16-byte byte array which are encoded big-endian for decimal128. */
+ byte[] getBinaryDecimal(int i) {
+ long valueBufferAddress = getValueVector().getDataBuffer().memoryAddress();
+ Platform.copyMemory(
+ null,
+ valueBufferAddress + (long) i * DECIMAL_BYTE_WIDTH,
+ DECIMAL_BYTES,
+ Platform.BYTE_ARRAY_OFFSET,
+ DECIMAL_BYTE_WIDTH);
+ // Decimal is stored little-endian in Arrow, so we need to reverse the bytes here
+ for (int j = 0, k = DECIMAL_BYTE_WIDTH - 1; j < DECIMAL_BYTE_WIDTH / 2; j++, k--) {
+ byte tmp = DECIMAL_BYTES[j];
+ DECIMAL_BYTES[j] = DECIMAL_BYTES[k];
+ DECIMAL_BYTES[k] = tmp;
+ }
+ return DECIMAL_BYTES;
+ }
+
+ @Override
+ public boolean getBoolean(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public byte getByte(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public short getShort(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public int getInt(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public long getLong(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public float getFloat(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public double getDouble(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public UTF8String getUTF8String(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public byte[] getBinary(int rowId) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public ColumnarArray getArray(int i) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public ColumnarMap getMap(int i) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public ColumnVector getChild(int i) {
+ throw new UnsupportedOperationException("Not yet supported");
+ }
+
+ @Override
+ public void close() {
+ getValueVector().close();
+ }
+
+ DictionaryProvider getDictionaryProvider() {
+ throw new UnsupportedOperationException("Not implemented");
+ }
+
+ abstract ValueVector getValueVector();
+
+ /**
+ * Returns a zero-copying new vector that contains the values from [offset, offset + length).
+ *
+ * @param offset the offset of the new vector
+ * @param length the length of the new vector
+ * @return the new vector
+ */
+ public abstract CometVector slice(int offset, int length);
+
+ /**
+ * Returns a corresponding `CometVector` implementation based on the given Arrow `ValueVector`.
+ *
+ * @param vector Arrow `ValueVector`
+ * @param useDecimal128 Whether to use Decimal128 for decimal column
+ * @return `CometVector` implementation
+ */
+ protected static CometVector getVector(
+ ValueVector vector, boolean useDecimal128, DictionaryProvider dictionaryProvider) {
+ if (vector instanceof StructVector) {
+ return new CometStructVector(vector, useDecimal128);
+ } else if (vector instanceof ListVector) {
+ return new CometListVector(vector, useDecimal128);
+ } else {
+ DictionaryEncoding dictionaryEncoding = vector.getField().getDictionary();
+ CometPlainVector cometVector = new CometPlainVector(vector, useDecimal128);
+
+ if (dictionaryEncoding == null) {
+ return cometVector;
+ } else {
+ Dictionary dictionary = dictionaryProvider.lookup(dictionaryEncoding.getId());
+ CometPlainVector dictionaryVector =
+ new CometPlainVector(dictionary.getVector(), useDecimal128);
+ CometDictionary cometDictionary = new CometDictionary(dictionaryVector);
+
+ return new CometDictionaryVector(
+ cometVector, cometDictionary, dictionaryProvider, useDecimal128);
+ }
+ }
+ }
+
+ protected static CometVector getVector(ValueVector vector, boolean useDecimal128) {
+ return getVector(vector, useDecimal128, null);
+ }
+}
diff --git a/common/src/main/resources/log4j2.properties b/common/src/main/resources/log4j2.properties
new file mode 100644
index 000000000..04cdf8533
--- /dev/null
+++ b/common/src/main/resources/log4j2.properties
@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+
+# Set everything to be logged to the file target/unit-tests.log
+rootLogger.level = info
+rootLogger.appenderRef.file.ref = ${sys:test.appender:-File}
+
+appender.file.type = File
+appender.file.name = File
+appender.file.fileName = target/unit-tests.log
+appender.file.layout.type = PatternLayout
+appender.file.layout.pattern = %d{yy/MM/dd HH:mm:ss.SSS} %t %p %c{1}: %m%n
+
+# Tests that launch java subprocesses can set the "test.appender" system property to
+# "console" to avoid having the child process's logs overwrite the unit test's
+# log file.
+appender.console.type = Console
+appender.console.name = console
+appender.console.target = SYSTEM_ERR
+appender.console.layout.type = PatternLayout
+appender.console.layout.pattern = %t: %m%n
+
+# Ignore messages below warning level from Jetty, because it's a bit verbose
+logger.jetty.name = org.sparkproject.jetty
+logger.jetty.level = warn
+
diff --git a/common/src/main/scala/org/apache/comet/CometConf.scala b/common/src/main/scala/org/apache/comet/CometConf.scala
new file mode 100644
index 000000000..7f83d9296
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/CometConf.scala
@@ -0,0 +1,454 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet
+
+import java.util.concurrent.TimeUnit
+
+import org.apache.spark.network.util.ByteUnit
+import org.apache.spark.network.util.JavaUtils
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Configurations for a Comet application. Mostly inspired by [[SQLConf]] in Spark.
+ *
+ * To get the value of a Comet config key from a [[SQLConf]], you can do the following:
+ *
+ * {{{
+ * CometConf.COMET_ENABLED.get
+ * }}}
+ *
+ * which retrieves the config value from the thread-local [[SQLConf]] object. Alternatively, you
+ * can also explicitly pass a [[SQLConf]] object to the `get` method.
+ */
+object CometConf {
+ def conf(key: String): ConfigBuilder = ConfigBuilder(key)
+
+ val COMET_EXEC_CONFIG_PREFIX = "spark.comet.exec";
+
+ val COMET_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.enabled")
+ .doc(
+ "Whether to enable Comet extension for Spark. When this is turned on, Spark will use " +
+ "Comet to read Parquet data source. Note that to enable native vectorized execution, " +
+ "both this config and 'spark.comet.exec.enabled' need to be enabled. By default, this " +
+ "config is the value of the env var `ENABLE_COMET` if set, or true otherwise.")
+ .booleanConf
+ .createWithDefault(sys.env.getOrElse("ENABLE_COMET", "true").toBoolean)
+
+ val COMET_SCAN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.scan.enabled")
+ .doc(
+ "Whether to enable Comet scan. When this is turned on, Spark will use Comet to read " +
+ "Parquet data source. Note that to enable native vectorized execution, both this " +
+ "config and 'spark.comet.exec.enabled' need to be enabled. By default, this config " +
+ "is true.")
+ .booleanConf
+ .createWithDefault(true)
+
+ val COMET_EXEC_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.enabled")
+ .doc(
+ "Whether to enable Comet native vectorized execution for Spark. This controls whether " +
+ "Spark should convert operators into their Comet counterparts and execute them in " +
+ "native space. Note: each operator is associated with a separate config in the " +
+ "format of 'spark.comet.exec..enabled' at the moment, and both the " +
+ "config and this need to be turned on, in order for the operator to be executed in " +
+ "native. By default, this config is false.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_MEMORY_OVERHEAD: OptionalConfigEntry[Long] = conf("spark.comet.memoryOverhead")
+ .doc(
+ "The amount of additional memory to be allocated per executor process for Comet, in MiB. " +
+ "This config is optional. If this is not specified, it will be set to " +
+ "`spark.comet.memory.overhead.factor` * `spark.executor.memory`. " +
+ "This is memory that accounts for things like Comet native execution, etc.")
+ .bytesConf(ByteUnit.MiB)
+ .createOptional
+
+ val COMET_MEMORY_OVERHEAD_FACTOR: ConfigEntry[Double] = conf(
+ "spark.comet.memory.overhead.factor")
+ .doc(
+ "Fraction of executor memory to be allocated as additional non-heap memory per executor " +
+ "process for Comet. Default value is 0.2.")
+ .doubleConf
+ .checkValue(
+ factor => factor > 0,
+ "Ensure that Comet memory overhead factor is a double greater than 0")
+ .createWithDefault(0.2)
+
+ val COMET_MEMORY_OVERHEAD_MIN_MIB: ConfigEntry[Long] = conf("spark.comet.memory.overhead.min")
+ .doc("Minimum amount of additional memory to be allocated per executor process for Comet, " +
+ "in MiB.")
+ .bytesConf(ByteUnit.MiB)
+ .checkValue(
+ _ >= 0,
+ "Ensure that Comet memory overhead min is a long greater than or equal to 0")
+ .createWithDefault(384)
+
+ val COMET_EXEC_ALL_OPERATOR_ENABLED: ConfigEntry[Boolean] = conf(
+ s"$COMET_EXEC_CONFIG_PREFIX.all.enabled")
+ .doc(
+ "Whether to enable all Comet operators. By default, this config is false. Note that " +
+ "this config precedes all separate config 'spark.comet.exec..enabled'. " +
+ "That being said, if this config is enabled, separate configs are ignored.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_EXEC_ALL_EXPR_ENABLED: ConfigEntry[Boolean] =
+ conf(s"$COMET_EXEC_CONFIG_PREFIX.all.expr.enabled")
+ .doc(
+ "Whether to enable all Comet exprs. By default, this config is false. Note that " +
+ "this config precedes all separate config 'spark.comet.exec..enabled'. " +
+ "That being said, if this config is enabled, separate configs are ignored.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_DEBUG_ENABLED: ConfigEntry[Boolean] =
+ conf("spark.comet.debug.enabled")
+ .doc(
+ "Whether to enable debug mode for Comet. By default, this config is false. " +
+ "When enabled, Comet will do additional checks for debugging purpose. For example, " +
+ "validating array when importing arrays from JVM at native side. Note that these " +
+ "checks may be expensive in performance and should only be enabled for debugging " +
+ "purpose.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_BATCH_SIZE: ConfigEntry[Int] = conf("spark.comet.batchSize")
+ .doc("The columnar batch size, i.e., the maximum number of rows that a batch can contain.")
+ .intConf
+ .createWithDefault(8192)
+
+ val COMET_EXEC_MEMORY_FRACTION: ConfigEntry[Double] = conf("spark.comet.exec.memoryFraction")
+ .doc(
+ "The fraction of memory from Comet memory overhead that the native memory " +
+ "manager can use for execution. The purpose of this config is to set aside memory for " +
+ "untracked data structures, as well as imprecise size estimation during memory " +
+ "acquisition. Default value is 0.7.")
+ .doubleConf
+ .createWithDefault(0.7)
+
+ val COMET_PARQUET_ENABLE_DIRECT_BUFFER: ConfigEntry[Boolean] = conf(
+ "spark.comet.parquet.enable.directBuffer")
+ .doc("Whether to use Java direct byte buffer when reading Parquet. By default, this is false")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_SCAN_PREFETCH_ENABLED: ConfigEntry[Boolean] =
+ conf("spark.comet.scan.preFetch.enabled")
+ .doc("Whether to enable pre-fetching feature of CometScan. By default is disabled.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_SCAN_PREFETCH_THREAD_NUM: ConfigEntry[Int] =
+ conf("spark.comet.scan.preFetch.threadNum")
+ .doc(
+ "The number of threads running pre-fetching for CometScan. Effective if " +
+ s"${COMET_SCAN_PREFETCH_ENABLED.key} is enabled. By default it is 2. Note that more " +
+ "pre-fetching threads means more memory requirement to store pre-fetched row groups.")
+ .intConf
+ .createWithDefault(2)
+
+ val COMET_NATIVE_LOAD_REQUIRED: ConfigEntry[Boolean] = conf("spark.comet.nativeLoadRequired")
+ .doc(
+ "Whether to require Comet native library to load successfully when Comet is enabled. " +
+ "If not, Comet will silently fallback to Spark when it fails to load the native lib. " +
+ "Otherwise, an error will be thrown and the Spark job will be aborted.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_EXCEPTION_ON_LEGACY_DATE_TIMESTAMP: ConfigEntry[Boolean] =
+ conf("spark.comet.exceptionOnDatetimeRebase")
+ .doc("Whether to throw exception when seeing dates/timestamps from the legacy hybrid " +
+ "(Julian + Gregorian) calendar. Since Spark 3, dates/timestamps were written according " +
+ "to the Proleptic Gregorian calendar. When this is true, Comet will " +
+ "throw exceptions when seeing these dates/timestamps that were written by Spark version " +
+ "before 3.0. If this is false, these dates/timestamps will be read as if they were " +
+ "written to the Proleptic Gregorian calendar and will not be rebased.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_USE_DECIMAL_128: ConfigEntry[Boolean] = conf("spark.comet.use.decimal128")
+ .internal()
+ .doc("If true, Comet will always use 128 bits to represent a decimal value, regardless of " +
+ "its precision. If false, Comet will use 32, 64 and 128 bits respectively depending on " +
+ "the precision. N.B. this is NOT a user-facing config but should be inferred and set by " +
+ "Comet itself.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_USE_LAZY_MATERIALIZATION: ConfigEntry[Boolean] = conf(
+ "spark.comet.use.lazyMaterialization")
+ .internal()
+ .doc(
+ "Whether to enable lazy materialization for Comet. When this is turned on, Comet will " +
+ "read Parquet data source lazily for string and binary columns. For filter operations, " +
+ "lazy materialization will improve read performance by skipping unused pages.")
+ .booleanConf
+ .createWithDefault(true)
+
+ val COMET_SCHEMA_EVOLUTION_ENABLED: ConfigEntry[Boolean] = conf(
+ "spark.comet.schemaEvolution.enabled")
+ .internal()
+ .doc(
+ "Whether to enable schema evolution in Comet. For instance, promoting a integer " +
+ "column to a long column, a float column to a double column, etc. This is automatically" +
+ "enabled when reading from Iceberg tables.")
+ .booleanConf
+ .createWithDefault(false)
+}
+
+object ConfigHelpers {
+ def toNumber[T](s: String, converter: String => T, key: String, configType: String): T = {
+ try {
+ converter(s.trim)
+ } catch {
+ case _: NumberFormatException =>
+ throw new IllegalArgumentException(s"$key should be $configType, but was $s")
+ }
+ }
+
+ def toBoolean(s: String, key: String): Boolean = {
+ try {
+ s.trim.toBoolean
+ } catch {
+ case _: IllegalArgumentException =>
+ throw new IllegalArgumentException(s"$key should be boolean, but was $s")
+ }
+ }
+
+ def stringToSeq[T](str: String, converter: String => T): Seq[T] = {
+ Utils.stringToSeq(str).map(converter)
+ }
+
+ def seqToString[T](v: Seq[T], stringConverter: T => String): String = {
+ v.map(stringConverter).mkString(",")
+ }
+
+ def timeFromString(str: String, unit: TimeUnit): Long = JavaUtils.timeStringAs(str, unit)
+
+ def timeToString(v: Long, unit: TimeUnit): String =
+ TimeUnit.MILLISECONDS.convert(v, unit) + "ms"
+
+ def byteFromString(str: String, unit: ByteUnit): Long = {
+ val (input, multiplier) =
+ if (str.nonEmpty && str.charAt(0) == '-') {
+ (str.substring(1), -1)
+ } else {
+ (str, 1)
+ }
+ multiplier * JavaUtils.byteStringAs(input, unit)
+ }
+
+ def byteToString(v: Long, unit: ByteUnit): String = unit.convertTo(v, ByteUnit.BYTE) + "b"
+}
+
+private class TypedConfigBuilder[T](
+ val parent: ConfigBuilder,
+ val converter: String => T,
+ val stringConverter: T => String) {
+
+ import ConfigHelpers._
+
+ def this(parent: ConfigBuilder, converter: String => T) = {
+ this(parent, converter, Option(_).map(_.toString).orNull)
+ }
+
+ /** Apply a transformation to the user-provided values of the config entry. */
+ def transform(fn: T => T): TypedConfigBuilder[T] = {
+ new TypedConfigBuilder(parent, s => fn(converter(s)), stringConverter)
+ }
+
+ /** Checks if the user-provided value for the config matches the validator. */
+ def checkValue(validator: T => Boolean, errorMsg: String): TypedConfigBuilder[T] = {
+ transform { v =>
+ if (!validator(v)) {
+ throw new IllegalArgumentException(s"'$v' in ${parent.key} is invalid. $errorMsg")
+ }
+ v
+ }
+ }
+
+ /** Check that user-provided values for the config match a pre-defined set. */
+ def checkValues(validValues: Set[T]): TypedConfigBuilder[T] = {
+ transform { v =>
+ if (!validValues.contains(v)) {
+ throw new IllegalArgumentException(
+ s"The value of ${parent.key} should be one of ${validValues.mkString(", ")}, but was $v")
+ }
+ v
+ }
+ }
+
+ /** Turns the config entry into a sequence of values of the underlying type. */
+ def toSequence: TypedConfigBuilder[Seq[T]] = {
+ new TypedConfigBuilder(parent, stringToSeq(_, converter), seqToString(_, stringConverter))
+ }
+
+ /** Creates a [[ConfigEntry]] that does not have a default value. */
+ def createOptional: OptionalConfigEntry[T] = {
+ new OptionalConfigEntry[T](
+ parent.key,
+ converter,
+ stringConverter,
+ parent._doc,
+ parent._public,
+ parent._version)
+ }
+
+ /** Creates a [[ConfigEntry]] that has a default value. */
+ def createWithDefault(default: T): ConfigEntry[T] = {
+ val transformedDefault = converter(stringConverter(default))
+ new ConfigEntryWithDefault[T](
+ parent.key,
+ transformedDefault,
+ converter,
+ stringConverter,
+ parent._doc,
+ parent._public,
+ parent._version)
+ }
+}
+
+private[comet] abstract class ConfigEntry[T](
+ val key: String,
+ val valueConverter: String => T,
+ val stringConverter: T => String,
+ val doc: String,
+ val isPublic: Boolean,
+ val version: String) {
+
+ /**
+ * Retrieves the config value from the given [[SQLConf]].
+ */
+ def get(conf: SQLConf): T
+
+ /**
+ * Retrieves the config value from the current thread-local [[SQLConf]]
+ * @return
+ */
+ def get(): T = get(SQLConf.get)
+
+ def defaultValue: Option[T] = None
+ def defaultValueString: String
+
+ override def toString: String = {
+ s"ConfigEntry(key=$key, defaultValue=$defaultValueString, doc=$doc, " +
+ s"public=$isPublic, version=$version)"
+ }
+}
+
+private[comet] class ConfigEntryWithDefault[T](
+ key: String,
+ _defaultValue: T,
+ valueConverter: String => T,
+ stringConverter: T => String,
+ doc: String,
+ isPublic: Boolean,
+ version: String)
+ extends ConfigEntry(key, valueConverter, stringConverter, doc, isPublic, version) {
+ override def defaultValue: Option[T] = Some(_defaultValue)
+ override def defaultValueString: String = stringConverter(_defaultValue)
+
+ def get(conf: SQLConf): T = {
+ val tmp = conf.getConfString(key, null)
+ if (tmp == null) {
+ _defaultValue
+ } else {
+ valueConverter(tmp)
+ }
+ }
+}
+
+private[comet] class OptionalConfigEntry[T](
+ key: String,
+ val rawValueConverter: String => T,
+ val rawStringConverter: T => String,
+ doc: String,
+ isPublic: Boolean,
+ version: String)
+ extends ConfigEntry[Option[T]](
+ key,
+ s => Some(rawValueConverter(s)),
+ v => v.map(rawStringConverter).orNull,
+ doc,
+ isPublic,
+ version) {
+
+ override def defaultValueString: String = ConfigEntry.UNDEFINED
+
+ override def get(conf: SQLConf): Option[T] = {
+ Option(conf.getConfString(key, null)).map(rawValueConverter)
+ }
+}
+
+private[comet] case class ConfigBuilder(key: String) {
+ import ConfigHelpers._
+
+ var _public = true
+ var _doc = ""
+ var _version = ""
+
+ def internal(): ConfigBuilder = {
+ _public = false
+ this
+ }
+
+ def doc(s: String): ConfigBuilder = {
+ _doc = s
+ this
+ }
+
+ def version(v: String): ConfigBuilder = {
+ _version = v
+ this
+ }
+
+ def intConf: TypedConfigBuilder[Int] = {
+ new TypedConfigBuilder(this, toNumber(_, _.toInt, key, "int"))
+ }
+
+ def longConf: TypedConfigBuilder[Long] = {
+ new TypedConfigBuilder(this, toNumber(_, _.toLong, key, "long"))
+ }
+
+ def doubleConf: TypedConfigBuilder[Double] = {
+ new TypedConfigBuilder(this, toNumber(_, _.toDouble, key, "double"))
+ }
+
+ def booleanConf: TypedConfigBuilder[Boolean] = {
+ new TypedConfigBuilder(this, toBoolean(_, key))
+ }
+
+ def stringConf: TypedConfigBuilder[String] = {
+ new TypedConfigBuilder(this, v => v)
+ }
+
+ def timeConf(unit: TimeUnit): TypedConfigBuilder[Long] = {
+ new TypedConfigBuilder(this, timeFromString(_, unit), timeToString(_, unit))
+ }
+
+ def bytesConf(unit: ByteUnit): TypedConfigBuilder[Long] = {
+ new TypedConfigBuilder(this, byteFromString(_, unit), byteToString(_, unit))
+ }
+}
+
+private object ConfigEntry {
+ val UNDEFINED = ""
+}
diff --git a/common/src/main/scala/org/apache/comet/Constants.scala b/common/src/main/scala/org/apache/comet/Constants.scala
new file mode 100644
index 000000000..83b570fc3
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/Constants.scala
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet
+
+object Constants {
+ val LOG_CONF_PATH = "comet.log.file.path"
+ val LOG_CONF_NAME = "log4rs.yaml"
+}
diff --git a/common/src/main/scala/org/apache/comet/package.scala b/common/src/main/scala/org/apache/comet/package.scala
new file mode 100644
index 000000000..c9aca7538
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/package.scala
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache
+
+import java.util.Properties
+
+package object comet {
+
+ /**
+ * Provides access to build information about the Comet libraries. This will be used by the
+ * benchmarking software to provide the source revision and repository. In addition, the build
+ * information is included to aid in future debugging efforts for releases.
+ */
+ private object CometBuildInfo {
+
+ val (
+ cometVersion: String,
+ cometBranch: String,
+ cometRevision: String,
+ cometBuildUserName: String,
+ cometBuildUserEmail: String,
+ cometRepoUrl: String,
+ cometBuildTimestamp: String) = {
+ val resourceStream = Thread
+ .currentThread()
+ .getContextClassLoader
+ .getResourceAsStream("comet-git-info.properties")
+ if (resourceStream == null) {
+ throw new CometRuntimeException("Could not find comet-git-info.properties")
+ }
+
+ try {
+ val unknownProp = ""
+ val props = new Properties()
+ props.load(resourceStream)
+ (
+ props.getProperty("git.build.version", unknownProp),
+ props.getProperty("git.branch", unknownProp),
+ props.getProperty("git.commit.id.full", unknownProp),
+ props.getProperty("git.build.user.name", unknownProp),
+ props.getProperty("git.build.user.email", unknownProp),
+ props.getProperty("git.remote.origin.url", unknownProp),
+ props.getProperty("git.build.time", unknownProp))
+ } catch {
+ case e: Exception =>
+ throw new CometRuntimeException(
+ "Error loading properties from comet-git-info.properties",
+ e)
+ } finally {
+ if (resourceStream != null) {
+ try {
+ resourceStream.close()
+ } catch {
+ case e: Exception =>
+ throw new CometRuntimeException("Error closing Comet build info resource stream", e)
+ }
+ }
+ }
+ }
+ }
+
+ val COMET_VERSION = CometBuildInfo.cometVersion
+ val COMET_BRANCH = CometBuildInfo.cometBranch
+ val COMET_REVISION = CometBuildInfo.cometRevision
+ val COMET_BUILD_USER_EMAIL = CometBuildInfo.cometBuildUserEmail
+ val COMET_BUILD_USER_NAME = CometBuildInfo.cometBuildUserName
+ val COMET_REPO_URL = CometBuildInfo.cometRepoUrl
+ val COMET_BUILD_TIMESTAMP = CometBuildInfo.cometBuildTimestamp
+
+}
diff --git a/common/src/main/scala/org/apache/comet/parquet/CometParquetUtils.scala b/common/src/main/scala/org/apache/comet/parquet/CometParquetUtils.scala
new file mode 100644
index 000000000..d851067b5
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/parquet/CometParquetUtils.scala
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+
+object CometParquetUtils {
+ private val PARQUET_FIELD_ID_WRITE_ENABLED = "spark.sql.parquet.fieldId.write.enabled"
+ private val PARQUET_FIELD_ID_READ_ENABLED = "spark.sql.parquet.fieldId.read.enabled"
+ private val IGNORE_MISSING_PARQUET_FIELD_ID = "spark.sql.parquet.fieldId.read.ignoreMissing"
+
+ def writeFieldId(conf: SQLConf): Boolean =
+ conf.getConfString(PARQUET_FIELD_ID_WRITE_ENABLED, "false").toBoolean
+
+ def writeFieldId(conf: Configuration): Boolean =
+ conf.getBoolean(PARQUET_FIELD_ID_WRITE_ENABLED, false)
+
+ def readFieldId(conf: SQLConf): Boolean =
+ conf.getConfString(PARQUET_FIELD_ID_READ_ENABLED, "false").toBoolean
+
+ def ignoreMissingIds(conf: SQLConf): Boolean =
+ conf.getConfString(IGNORE_MISSING_PARQUET_FIELD_ID, "false").toBoolean
+
+ // The following is copied from QueryExecutionErrors
+ // TODO: remove after dropping Spark 3.2.0 support and directly use
+ // QueryExecutionErrors.foundDuplicateFieldInFieldIdLookupModeError
+ def foundDuplicateFieldInFieldIdLookupModeError(
+ requiredId: Int,
+ matchedFields: String): Throwable = {
+ new RuntimeException(s"""
+ |Found duplicate field(s) "$requiredId": $matchedFields
+ |in id mapping mode
+ """.stripMargin.replaceAll("\n", " "))
+ }
+
+ // The followings are copied from org.apache.spark.sql.execution.datasources.parquet.ParquetUtils
+ // TODO: remove after dropping Spark 3.2.0 support and directly use ParquetUtils
+ /**
+ * A StructField metadata key used to set the field id of a column in the Parquet schema.
+ */
+ val FIELD_ID_METADATA_KEY = "parquet.field.id"
+
+ /**
+ * Whether there exists a field in the schema, whether inner or leaf, has the parquet field ID
+ * metadata.
+ */
+ def hasFieldIds(schema: StructType): Boolean = {
+ def recursiveCheck(schema: DataType): Boolean = {
+ schema match {
+ case st: StructType =>
+ st.exists(field => hasFieldId(field) || recursiveCheck(field.dataType))
+
+ case at: ArrayType => recursiveCheck(at.elementType)
+
+ case mt: MapType => recursiveCheck(mt.keyType) || recursiveCheck(mt.valueType)
+
+ case _ =>
+ // No need to really check primitive types, just to terminate the recursion
+ false
+ }
+ }
+ if (schema.isEmpty) false else recursiveCheck(schema)
+ }
+
+ def hasFieldId(field: StructField): Boolean =
+ field.metadata.contains(FIELD_ID_METADATA_KEY)
+
+ def getFieldId(field: StructField): Int = {
+ require(
+ hasFieldId(field),
+ s"The key `$FIELD_ID_METADATA_KEY` doesn't exist in the metadata of " + field)
+ try {
+ Math.toIntExact(field.metadata.getLong(FIELD_ID_METADATA_KEY))
+ } catch {
+ case _: ArithmeticException | _: ClassCastException =>
+ throw new IllegalArgumentException(
+ s"The key `$FIELD_ID_METADATA_KEY` must be a 32-bit integer")
+ }
+ }
+}
diff --git a/common/src/main/scala/org/apache/comet/parquet/CometReaderThreadPool.scala b/common/src/main/scala/org/apache/comet/parquet/CometReaderThreadPool.scala
new file mode 100644
index 000000000..ca13bba0c
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/parquet/CometReaderThreadPool.scala
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet
+
+import java.util.concurrent.{Executors, ExecutorService, ThreadFactory}
+import java.util.concurrent.atomic.AtomicLong
+
+abstract class CometReaderThreadPool {
+ private var threadPool: Option[ExecutorService] = None
+
+ protected def threadNamePrefix: String
+
+ private def initThreadPool(maxThreads: Int): ExecutorService = synchronized {
+ if (threadPool.isEmpty) {
+ val threadFactory: ThreadFactory = new ThreadFactory() {
+ private val defaultThreadFactory = Executors.defaultThreadFactory
+ val count = new AtomicLong(0)
+
+ override def newThread(r: Runnable): Thread = {
+ val thread = defaultThreadFactory.newThread(r)
+ thread.setName(s"${threadNamePrefix}_${count.getAndIncrement()}")
+ thread.setDaemon(true)
+ thread
+ }
+ }
+
+ val threadPoolExecutor = Executors.newFixedThreadPool(maxThreads, threadFactory)
+ threadPool = Some(threadPoolExecutor)
+ }
+
+ threadPool.get
+ }
+
+ def getOrCreateThreadPool(numThreads: Int): ExecutorService = {
+ threadPool.getOrElse(initThreadPool(numThreads))
+ }
+
+}
+
+// A thread pool used for pre-fetching files.
+object CometPrefetchThreadPool extends CometReaderThreadPool {
+ override def threadNamePrefix: String = "prefetch_thread"
+}
+
+// Thread pool used by the Parquet parallel reader
+object CometFileReaderThreadPool extends CometReaderThreadPool {
+ override def threadNamePrefix: String = "file_reader_thread"
+}
diff --git a/common/src/main/scala/org/apache/comet/shims/ShimBatchReader.scala b/common/src/main/scala/org/apache/comet/shims/ShimBatchReader.scala
new file mode 100644
index 000000000..ece4cfbe5
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/shims/ShimBatchReader.scala
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.shims
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.execution.datasources.PartitionedFile
+
+object ShimBatchReader {
+
+ // TODO: remove after dropping Spark 3.2 & 3.3 support and directly call PartitionedFile
+ def newPartitionedFile(partitionValues: InternalRow, file: String): PartitionedFile =
+ classOf[PartitionedFile].getDeclaredConstructors
+ .map(c =>
+ c.getParameterCount match {
+ case 5 =>
+ c.newInstance(
+ partitionValues,
+ file,
+ Long.box(-1), // -1 means we read the entire file
+ Long.box(-1),
+ Array.empty[String])
+ case 7 =>
+ c.newInstance(
+ partitionValues,
+ c.getParameterTypes()(1)
+ .getConstructor(classOf[String])
+ .newInstance(file)
+ .asInstanceOf[AnyRef],
+ Long.box(-1), // -1 means we read the entire file
+ Long.box(-1),
+ Array.empty[String],
+ Long.box(0),
+ Long.box(0))
+ })
+ .head
+ .asInstanceOf[PartitionedFile]
+}
diff --git a/common/src/main/scala/org/apache/comet/shims/ShimFileFormat.scala b/common/src/main/scala/org/apache/comet/shims/ShimFileFormat.scala
new file mode 100644
index 000000000..5ab7eaf4f
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/shims/ShimFileFormat.scala
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.shims
+
+object ShimFileFormat {
+
+ // TODO: remove after dropping Spark 3.2 & 3.3 support and directly use FileFormat.ROW_INDEX
+ val ROW_INDEX = "row_index"
+
+ // A name for a temporary column that holds row indexes computed by the file format reader
+ // until they can be placed in the _metadata struct.
+ // TODO: remove after dropping Spark 3.2 & 3.3 support and directly use
+ // FileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME
+ val ROW_INDEX_TEMPORARY_COLUMN_NAME: String = s"_tmp_metadata_$ROW_INDEX"
+
+ // TODO: remove after dropping Spark 3.2 support and use FileFormat.OPTION_RETURNING_BATCH
+ val OPTION_RETURNING_BATCH = "returning_batch"
+}
diff --git a/common/src/main/scala/org/apache/comet/shims/ShimResolveDefaultColumns.scala b/common/src/main/scala/org/apache/comet/shims/ShimResolveDefaultColumns.scala
new file mode 100644
index 000000000..8a30c8e00
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/shims/ShimResolveDefaultColumns.scala
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.shims
+
+import scala.util.Try
+
+import org.apache.spark.sql.types.{StructField, StructType}
+
+object ShimResolveDefaultColumns {
+ // TODO: remove after dropping Spark 3.2 & 3.3 support and directly use ResolveDefaultColumns
+ def getExistenceDefaultValue(field: StructField): Any =
+ Try {
+ // scalastyle:off classforname
+ Class.forName("org.apache.spark.sql.catalyst.util.ResolveDefaultColumns$")
+ // scalastyle:on classforname
+ }.map { objClass =>
+ val objInstance = objClass.getField("MODULE$").get(null)
+ val method = objClass.getMethod("getExistenceDefaultValues", classOf[StructType])
+ method.invoke(objInstance, StructType(Seq(field))).asInstanceOf[Array[Any]].head
+ }.getOrElse(null)
+}
diff --git a/common/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/common/src/main/scala/org/apache/comet/vector/NativeUtil.scala
new file mode 100644
index 000000000..1e27ed8f0
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/vector/NativeUtil.scala
@@ -0,0 +1,141 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector
+
+import scala.collection.mutable
+
+import org.apache.arrow.c.{ArrowArray, ArrowImporter, ArrowSchema, CDataDictionaryProvider, Data}
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector._
+import org.apache.spark.SparkException
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+class NativeUtil {
+ private val allocator = new RootAllocator(Long.MaxValue)
+ private val dictionaryProvider: CDataDictionaryProvider = new CDataDictionaryProvider
+ private val importer = new ArrowImporter(allocator)
+
+ /**
+ * Exports a Comet `ColumnarBatch` into a list of memory addresses that can be consumed by the
+ * native execution.
+ *
+ * @param batch
+ * the input Comet columnar batch
+ * @return
+ * a list containing pairs of memory addresses in the format of (address of Arrow array,
+ * address of Arrow schema)
+ */
+ def exportBatch(batch: ColumnarBatch): Array[Long] = {
+ val vectors = (0 until batch.numCols()).flatMap { index =>
+ batch.column(index) match {
+ case a: CometVector =>
+ val valueVector = a.getValueVector
+
+ val provider = if (valueVector.getField.getDictionary != null) {
+ a.getDictionaryProvider
+ } else {
+ null
+ }
+
+ val arrowSchema = ArrowSchema.allocateNew(allocator)
+ val arrowArray = ArrowArray.allocateNew(allocator)
+ Data.exportVector(
+ allocator,
+ getFieldVector(valueVector),
+ provider,
+ arrowArray,
+ arrowSchema)
+
+ Seq((arrowArray, arrowSchema))
+ case c =>
+ throw new SparkException(
+ "Comet execution only takes Arrow Arrays, but got " +
+ s"${c.getClass}")
+ }
+ }
+
+ vectors.flatMap { pair =>
+ Seq(pair._1.memoryAddress(), pair._2.memoryAddress())
+ }.toArray
+ }
+
+ /**
+ * Imports a list of Arrow addresses from native execution, and return a list of Comet vectors.
+ *
+ * @param arrayAddress
+ * a list containing paris of Arrow addresses from the native, in the format of (address of
+ * Arrow array, address of Arrow schema)
+ * @return
+ * a list of Comet vectors
+ */
+ def importVector(arrayAddress: Array[Long]): Seq[CometVector] = {
+ val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
+
+ for (i <- arrayAddress.indices by 2) {
+ val arrowSchema = ArrowSchema.wrap(arrayAddress(i + 1))
+ val arrowArray = ArrowArray.wrap(arrayAddress(i))
+
+ // Native execution should always have 'useDecimal128' set to true since it doesn't support
+ // other cases.
+ arrayVectors += CometVector.getVector(
+ importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
+ true,
+ dictionaryProvider)
+
+ arrowArray.close()
+ arrowSchema.close()
+ }
+ arrayVectors.toSeq
+ }
+
+ /**
+ * Takes zero-copy slices of the input batch with given start index and maximum number of rows.
+ *
+ * @param batch
+ * Input batch
+ * @param startIndex
+ * Start index of the slice
+ * @param maxNumRows
+ * Maximum number of rows in the slice
+ * @return
+ * A new batch with the sliced vectors
+ */
+ def takeRows(batch: ColumnarBatch, startIndex: Int, maxNumRows: Int): ColumnarBatch = {
+ val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
+
+ for (i <- 0 until batch.numCols()) {
+ val column = batch.column(i).asInstanceOf[CometVector]
+ arrayVectors += column.slice(startIndex, maxNumRows)
+ }
+
+ new ColumnarBatch(arrayVectors.toArray, maxNumRows)
+ }
+
+ private def getFieldVector(valueVector: ValueVector): FieldVector = {
+ valueVector match {
+ case v @ (_: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector |
+ _: BigIntVector | _: Float4Vector | _: Float8Vector | _: VarCharVector |
+ _: DecimalVector | _: DateDayVector | _: TimeStampMicroTZVector | _: VarBinaryVector |
+ _: FixedSizeBinaryVector | _: TimeStampMicroVector) =>
+ v.asInstanceOf[FieldVector]
+ case _ => throw new SparkException(s"Unsupported Arrow Vector: ${valueVector.getClass}")
+ }
+ }
+}
diff --git a/common/src/main/scala/org/apache/comet/vector/StreamReader.scala b/common/src/main/scala/org/apache/comet/vector/StreamReader.scala
new file mode 100644
index 000000000..9c4f99602
--- /dev/null
+++ b/common/src/main/scala/org/apache/comet/vector/StreamReader.scala
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.vector
+
+import java.nio.channels.ReadableByteChannel
+
+import scala.collection.JavaConverters.collectionAsScalaIterableConverter
+
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel}
+import org.apache.arrow.vector.ipc.message.MessageChannelReader
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+/**
+ * A reader that consumes Arrow data from an input channel, and produces Comet batches.
+ */
+case class StreamReader(channel: ReadableByteChannel) extends AutoCloseable {
+ private var allocator = new RootAllocator(Long.MaxValue)
+ private val channelReader = new MessageChannelReader(new ReadChannel(channel), allocator)
+ private var arrowReader = new ArrowStreamReader(channelReader, allocator)
+ private var root = arrowReader.getVectorSchemaRoot
+
+ def nextBatch(): Option[ColumnarBatch] = {
+ if (arrowReader.loadNextBatch()) {
+ Some(rootAsBatch(root))
+ } else {
+ None
+ }
+ }
+
+ private def rootAsBatch(root: VectorSchemaRoot): ColumnarBatch = {
+ val columns = root.getFieldVectors.asScala.map { vector =>
+ // Native shuffle always uses decimal128.
+ CometVector.getVector(vector, true, arrowReader).asInstanceOf[ColumnVector]
+ }.toArray
+ val batch = new ColumnarBatch(columns)
+ batch.setNumRows(root.getRowCount)
+ batch
+ }
+
+ override def close(): Unit = {
+ if (root != null) {
+ arrowReader.close()
+ root.close()
+ allocator.close()
+
+ arrowReader = null
+ root = null
+ allocator = null
+ }
+ }
+}
diff --git a/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometParquetReadSupport.scala b/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometParquetReadSupport.scala
new file mode 100644
index 000000000..0e8a190c2
--- /dev/null
+++ b/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometParquetReadSupport.scala
@@ -0,0 +1,378 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.spark.sql.comet.parquet
+
+import java.util.{Locale, UUID}
+
+import scala.collection.JavaConverters._
+
+import org.apache.parquet.schema._
+import org.apache.parquet.schema.LogicalTypeAnnotation.ListLogicalTypeAnnotation
+import org.apache.parquet.schema.Type.Repetition
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.types._
+
+import org.apache.comet.parquet.CometParquetUtils
+
+/**
+ * This class is copied & slightly modified from [[ParquetReadSupport]] in Spark. Changes:
+ * - This doesn't extend from Parquet's `ReadSupport` class since that is used for row-based
+ * Parquet reader. Therefore, there is no `init`, `prepareForRead` as well as other methods
+ * that are unused.
+ */
+object CometParquetReadSupport {
+ val SPARK_PARQUET_SCHEMA_NAME = "spark_schema"
+
+ val EMPTY_MESSAGE: MessageType =
+ Types.buildMessage().named(SPARK_PARQUET_SCHEMA_NAME)
+
+ def generateFakeColumnName: String = s"_fake_name_${UUID.randomUUID()}"
+
+ def clipParquetSchema(
+ parquetSchema: MessageType,
+ catalystSchema: StructType,
+ caseSensitive: Boolean,
+ useFieldId: Boolean,
+ ignoreMissingIds: Boolean): MessageType = {
+ if (!ignoreMissingIds &&
+ !containsFieldIds(parquetSchema) &&
+ CometParquetUtils.hasFieldIds(catalystSchema)) {
+ throw new RuntimeException(
+ "Spark read schema expects field Ids, " +
+ "but Parquet file schema doesn't contain any field Ids.\n" +
+ "Please remove the field ids from Spark schema or ignore missing ids by " +
+ "setting `spark.sql.parquet.fieldId.read.ignoreMissing = true`\n" +
+ s"""
+ |Spark read schema:
+ |${catalystSchema.prettyJson}
+ |
+ |Parquet file schema:
+ |${parquetSchema.toString}
+ |""".stripMargin)
+ }
+ clipParquetSchema(parquetSchema, catalystSchema, caseSensitive, useFieldId)
+ }
+
+ /**
+ * Tailors `parquetSchema` according to `catalystSchema` by removing column paths don't exist in
+ * `catalystSchema`, and adding those only exist in `catalystSchema`.
+ */
+ def clipParquetSchema(
+ parquetSchema: MessageType,
+ catalystSchema: StructType,
+ caseSensitive: Boolean,
+ useFieldId: Boolean): MessageType = {
+ val clippedParquetFields = clipParquetGroupFields(
+ parquetSchema.asGroupType(),
+ catalystSchema,
+ caseSensitive,
+ useFieldId)
+ if (clippedParquetFields.isEmpty) {
+ EMPTY_MESSAGE
+ } else {
+ Types
+ .buildMessage()
+ .addFields(clippedParquetFields: _*)
+ .named(SPARK_PARQUET_SCHEMA_NAME)
+ }
+ }
+
+ private def clipParquetType(
+ parquetType: Type,
+ catalystType: DataType,
+ caseSensitive: Boolean,
+ useFieldId: Boolean): Type = {
+ val newParquetType = catalystType match {
+ case t: ArrayType if !isPrimitiveCatalystType(t.elementType) =>
+ // Only clips array types with nested type as element type.
+ clipParquetListType(parquetType.asGroupType(), t.elementType, caseSensitive, useFieldId)
+
+ case t: MapType
+ if !isPrimitiveCatalystType(t.keyType) ||
+ !isPrimitiveCatalystType(t.valueType) =>
+ // Only clips map types with nested key type or value type
+ clipParquetMapType(
+ parquetType.asGroupType(),
+ t.keyType,
+ t.valueType,
+ caseSensitive,
+ useFieldId)
+
+ case t: StructType =>
+ clipParquetGroup(parquetType.asGroupType(), t, caseSensitive, useFieldId)
+
+ case _ =>
+ // UDTs and primitive types are not clipped. For UDTs, a clipped version might not be able
+ // to be mapped to desired user-space types. So UDTs shouldn't participate schema merging.
+ parquetType
+ }
+
+ if (useFieldId && parquetType.getId != null) {
+ newParquetType.withId(parquetType.getId.intValue())
+ } else {
+ newParquetType
+ }
+ }
+
+ /**
+ * Whether a Catalyst [[DataType]] is primitive. Primitive [[DataType]] is not equivalent to
+ * [[AtomicType]]. For example, [[CalendarIntervalType]] is primitive, but it's not an
+ * [[AtomicType]].
+ */
+ private def isPrimitiveCatalystType(dataType: DataType): Boolean = {
+ dataType match {
+ case _: ArrayType | _: MapType | _: StructType => false
+ case _ => true
+ }
+ }
+
+ /**
+ * Clips a Parquet [[GroupType]] which corresponds to a Catalyst [[ArrayType]]. The element type
+ * of the [[ArrayType]] should also be a nested type, namely an [[ArrayType]], a [[MapType]], or
+ * a [[StructType]].
+ */
+ private def clipParquetListType(
+ parquetList: GroupType,
+ elementType: DataType,
+ caseSensitive: Boolean,
+ useFieldId: Boolean): Type = {
+ // Precondition of this method, should only be called for lists with nested element types.
+ assert(!isPrimitiveCatalystType(elementType))
+
+ // Unannotated repeated group should be interpreted as required list of required element, so
+ // list element type is just the group itself. Clip it.
+ if (parquetList.getLogicalTypeAnnotation == null &&
+ parquetList.isRepetition(Repetition.REPEATED)) {
+ clipParquetType(parquetList, elementType, caseSensitive, useFieldId)
+ } else {
+ assert(
+ parquetList.getLogicalTypeAnnotation.isInstanceOf[ListLogicalTypeAnnotation],
+ "Invalid Parquet schema. " +
+ "Logical type annotation of annotated Parquet lists must be ListLogicalTypeAnnotation: " +
+ parquetList.toString)
+
+ assert(
+ parquetList.getFieldCount == 1 && parquetList
+ .getType(0)
+ .isRepetition(Repetition.REPEATED),
+ "Invalid Parquet schema. " +
+ "LIST-annotated group should only have exactly one repeated field: " +
+ parquetList)
+
+ // Precondition of this method, should only be called for lists with nested element types.
+ assert(!parquetList.getType(0).isPrimitive)
+
+ val repeatedGroup = parquetList.getType(0).asGroupType()
+
+ // If the repeated field is a group with multiple fields, or the repeated field is a group
+ // with one field and is named either "array" or uses the LIST-annotated group's name with
+ // "_tuple" appended then the repeated type is the element type and elements are required.
+ // Build a new LIST-annotated group with clipped `repeatedGroup` as element type and the
+ // only field.
+ if (repeatedGroup.getFieldCount > 1 ||
+ repeatedGroup.getName == "array" ||
+ repeatedGroup.getName == parquetList.getName + "_tuple") {
+ Types
+ .buildGroup(parquetList.getRepetition)
+ .as(LogicalTypeAnnotation.listType())
+ .addField(clipParquetType(repeatedGroup, elementType, caseSensitive, useFieldId))
+ .named(parquetList.getName)
+ } else {
+ // Otherwise, the repeated field's type is the element type with the repeated field's
+ // repetition.
+ val newRepeatedGroup = Types
+ .repeatedGroup()
+ .addField(
+ clipParquetType(repeatedGroup.getType(0), elementType, caseSensitive, useFieldId))
+ .named(repeatedGroup.getName)
+
+ val newElementType = if (useFieldId && repeatedGroup.getId != null) {
+ newRepeatedGroup.withId(repeatedGroup.getId.intValue())
+ } else {
+ newRepeatedGroup
+ }
+
+ Types
+ .buildGroup(parquetList.getRepetition)
+ .as(LogicalTypeAnnotation.listType())
+ .addField(
+ Types
+ .repeatedGroup()
+ .addField(
+ clipParquetType(repeatedGroup.getType(0), elementType, caseSensitive, useFieldId))
+ .named(repeatedGroup.getName))
+ .addField(newElementType)
+ .named(parquetList.getName)
+ }
+ }
+ }
+
+ /**
+ * Clips a Parquet [[GroupType]] which corresponds to a Catalyst [[MapType]]. Either key type or
+ * value type of the [[MapType]] must be a nested type, namely an [[ArrayType]], a [[MapType]],
+ * or a [[StructType]].
+ */
+ private def clipParquetMapType(
+ parquetMap: GroupType,
+ keyType: DataType,
+ valueType: DataType,
+ caseSensitive: Boolean,
+ useFieldId: Boolean): GroupType = {
+ // Precondition of this method, only handles maps with nested key types or value types.
+ assert(!isPrimitiveCatalystType(keyType) || !isPrimitiveCatalystType(valueType))
+
+ val repeatedGroup = parquetMap.getType(0).asGroupType()
+ val parquetKeyType = repeatedGroup.getType(0)
+ val parquetValueType = repeatedGroup.getType(1)
+
+ val clippedRepeatedGroup = {
+ val newRepeatedGroup = Types
+ .repeatedGroup()
+ .as(repeatedGroup.getLogicalTypeAnnotation)
+ .addField(clipParquetType(parquetKeyType, keyType, caseSensitive, useFieldId))
+ .addField(clipParquetType(parquetValueType, valueType, caseSensitive, useFieldId))
+ .named(repeatedGroup.getName)
+ if (useFieldId && repeatedGroup.getId != null) {
+ newRepeatedGroup.withId(repeatedGroup.getId.intValue())
+ } else {
+ newRepeatedGroup
+ }
+ }
+
+ Types
+ .buildGroup(parquetMap.getRepetition)
+ .as(parquetMap.getLogicalTypeAnnotation)
+ .addField(clippedRepeatedGroup)
+ .named(parquetMap.getName)
+ }
+
+ /**
+ * Clips a Parquet [[GroupType]] which corresponds to a Catalyst [[StructType]].
+ *
+ * @return
+ * A clipped [[GroupType]], which has at least one field.
+ * @note
+ * Parquet doesn't allow creating empty [[GroupType]] instances except for empty
+ * [[MessageType]]. Because it's legal to construct an empty requested schema for column
+ * pruning.
+ */
+ private def clipParquetGroup(
+ parquetRecord: GroupType,
+ structType: StructType,
+ caseSensitive: Boolean,
+ useFieldId: Boolean): GroupType = {
+ val clippedParquetFields =
+ clipParquetGroupFields(parquetRecord, structType, caseSensitive, useFieldId)
+ Types
+ .buildGroup(parquetRecord.getRepetition)
+ .as(parquetRecord.getLogicalTypeAnnotation)
+ .addFields(clippedParquetFields: _*)
+ .named(parquetRecord.getName)
+ }
+
+ /**
+ * Clips a Parquet [[GroupType]] which corresponds to a Catalyst [[StructType]].
+ *
+ * @return
+ * A list of clipped [[GroupType]] fields, which can be empty.
+ */
+ private def clipParquetGroupFields(
+ parquetRecord: GroupType,
+ structType: StructType,
+ caseSensitive: Boolean,
+ useFieldId: Boolean): Seq[Type] = {
+ val toParquet = new CometSparkToParquetSchemaConverter(
+ writeLegacyParquetFormat = false,
+ useFieldId = useFieldId)
+ lazy val caseSensitiveParquetFieldMap =
+ parquetRecord.getFields.asScala.map(f => f.getName -> f).toMap
+ lazy val caseInsensitiveParquetFieldMap =
+ parquetRecord.getFields.asScala.groupBy(_.getName.toLowerCase(Locale.ROOT))
+ lazy val idToParquetFieldMap =
+ parquetRecord.getFields.asScala.filter(_.getId != null).groupBy(f => f.getId.intValue())
+
+ def matchCaseSensitiveField(f: StructField): Type = {
+ caseSensitiveParquetFieldMap
+ .get(f.name)
+ .map(clipParquetType(_, f.dataType, caseSensitive, useFieldId))
+ .getOrElse(toParquet.convertField(f))
+ }
+
+ def matchCaseInsensitiveField(f: StructField): Type = {
+ // Do case-insensitive resolution only if in case-insensitive mode
+ caseInsensitiveParquetFieldMap
+ .get(f.name.toLowerCase(Locale.ROOT))
+ .map { parquetTypes =>
+ if (parquetTypes.size > 1) {
+ // Need to fail if there is ambiguity, i.e. more than one field is matched
+ val parquetTypesString = parquetTypes.map(_.getName).mkString("[", ", ", "]")
+ throw QueryExecutionErrors.foundDuplicateFieldInCaseInsensitiveModeError(
+ f.name,
+ parquetTypesString)
+ } else {
+ clipParquetType(parquetTypes.head, f.dataType, caseSensitive, useFieldId)
+ }
+ }
+ .getOrElse(toParquet.convertField(f))
+ }
+
+ def matchIdField(f: StructField): Type = {
+ val fieldId = CometParquetUtils.getFieldId(f)
+ idToParquetFieldMap
+ .get(fieldId)
+ .map { parquetTypes =>
+ if (parquetTypes.size > 1) {
+ // Need to fail if there is ambiguity, i.e. more than one field is matched
+ val parquetTypesString = parquetTypes.map(_.getName).mkString("[", ", ", "]")
+ throw CometParquetUtils.foundDuplicateFieldInFieldIdLookupModeError(
+ fieldId,
+ parquetTypesString)
+ } else {
+ clipParquetType(parquetTypes.head, f.dataType, caseSensitive, useFieldId)
+ }
+ }
+ .getOrElse {
+ // When there is no ID match, we use a fake name to avoid a name match by accident
+ // We need this name to be unique as well, otherwise there will be type conflicts
+ toParquet.convertField(f.copy(name = generateFakeColumnName))
+ }
+ }
+
+ val shouldMatchById = useFieldId && CometParquetUtils.hasFieldIds(structType)
+ structType.map { f =>
+ if (shouldMatchById && CometParquetUtils.hasFieldId(f)) {
+ matchIdField(f)
+ } else if (caseSensitive) {
+ matchCaseSensitiveField(f)
+ } else {
+ matchCaseInsensitiveField(f)
+ }
+ }
+ }
+
+ /**
+ * Whether the parquet schema contains any field IDs.
+ */
+ private def containsFieldIds(schema: Type): Boolean = schema match {
+ case p: PrimitiveType => p.getId != null
+ // We don't require all fields to have IDs, so we use `exists` here.
+ case g: GroupType => g.getId != null || g.getFields.asScala.exists(containsFieldIds)
+ }
+}
diff --git a/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometSparkToParquetSchemaConverter.scala b/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometSparkToParquetSchemaConverter.scala
new file mode 100644
index 000000000..2c8187e18
--- /dev/null
+++ b/common/src/main/scala/org/apache/spark/sql/comet/parquet/CometSparkToParquetSchemaConverter.scala
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.spark.sql.comet.parquet
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.parquet.schema._
+import org.apache.parquet.schema.LogicalTypeAnnotation._
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName._
+import org.apache.parquet.schema.Type.Repetition._
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.execution.datasources.parquet.ParquetSchemaConverter
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+
+import org.apache.comet.parquet.CometParquetUtils
+
+/**
+ * This class is copied & modified from Spark's [[SparkToParquetSchemaConverter]] class.
+ */
+class CometSparkToParquetSchemaConverter(
+ writeLegacyParquetFormat: Boolean = SQLConf.PARQUET_WRITE_LEGACY_FORMAT.defaultValue.get,
+ outputTimestampType: SQLConf.ParquetOutputTimestampType.Value =
+ SQLConf.ParquetOutputTimestampType.INT96,
+ useFieldId: Boolean = CometParquetUtils.writeFieldId(new SQLConf)) {
+
+ def this(conf: SQLConf) = this(
+ writeLegacyParquetFormat = conf.writeLegacyParquetFormat,
+ outputTimestampType = conf.parquetOutputTimestampType,
+ useFieldId = CometParquetUtils.writeFieldId(conf))
+
+ def this(conf: Configuration) = this(
+ writeLegacyParquetFormat = conf.get(SQLConf.PARQUET_WRITE_LEGACY_FORMAT.key).toBoolean,
+ outputTimestampType = SQLConf.ParquetOutputTimestampType.withName(
+ conf.get(SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key)),
+ useFieldId = CometParquetUtils.writeFieldId(conf))
+
+ /**
+ * Converts a Spark SQL [[StructType]] to a Parquet [[MessageType]].
+ */
+ def convert(catalystSchema: StructType): MessageType = {
+ Types
+ .buildMessage()
+ .addFields(catalystSchema.map(convertField): _*)
+ .named(ParquetSchemaConverter.SPARK_PARQUET_SCHEMA_NAME)
+ }
+
+ /**
+ * Converts a Spark SQL [[StructField]] to a Parquet [[Type]].
+ */
+ def convertField(field: StructField): Type = {
+ val converted = convertField(field, if (field.nullable) OPTIONAL else REQUIRED)
+ if (useFieldId && CometParquetUtils.hasFieldId(field)) {
+ converted.withId(CometParquetUtils.getFieldId(field))
+ } else {
+ converted
+ }
+ }
+
+ private def convertField(field: StructField, repetition: Type.Repetition): Type = {
+
+ field.dataType match {
+ // ===================
+ // Simple atomic types
+ // ===================
+
+ case BooleanType =>
+ Types.primitive(BOOLEAN, repetition).named(field.name)
+
+ case ByteType =>
+ Types
+ .primitive(INT32, repetition)
+ .as(LogicalTypeAnnotation.intType(8, true))
+ .named(field.name)
+
+ case ShortType =>
+ Types
+ .primitive(INT32, repetition)
+ .as(LogicalTypeAnnotation.intType(16, true))
+ .named(field.name)
+
+ case IntegerType =>
+ Types.primitive(INT32, repetition).named(field.name)
+
+ case LongType =>
+ Types.primitive(INT64, repetition).named(field.name)
+
+ case FloatType =>
+ Types.primitive(FLOAT, repetition).named(field.name)
+
+ case DoubleType =>
+ Types.primitive(DOUBLE, repetition).named(field.name)
+
+ case StringType =>
+ Types
+ .primitive(BINARY, repetition)
+ .as(LogicalTypeAnnotation.stringType())
+ .named(field.name)
+
+ case DateType =>
+ Types
+ .primitive(INT32, repetition)
+ .as(LogicalTypeAnnotation.dateType())
+ .named(field.name)
+
+ // NOTE: Spark SQL can write timestamp values to Parquet using INT96, TIMESTAMP_MICROS or
+ // TIMESTAMP_MILLIS. TIMESTAMP_MICROS is recommended but INT96 is the default to keep the
+ // behavior same as before.
+ //
+ // As stated in PARQUET-323, Parquet `INT96` was originally introduced to represent nanosecond
+ // timestamp in Impala for some historical reasons. It's not recommended to be used for any
+ // other types and will probably be deprecated in some future version of parquet-format spec.
+ // That's the reason why parquet-format spec only defines `TIMESTAMP_MILLIS` and
+ // `TIMESTAMP_MICROS` which are both logical types annotating `INT64`.
+ //
+ // Originally, Spark SQL uses the same nanosecond timestamp type as Impala and Hive. Starting
+ // from Spark 1.5.0, we resort to a timestamp type with microsecond precision so that we can
+ // store a timestamp into a `Long`. This design decision is subject to change though, for
+ // example, we may resort to nanosecond precision in the future.
+ case TimestampType =>
+ outputTimestampType match {
+ case SQLConf.ParquetOutputTimestampType.INT96 =>
+ Types.primitive(INT96, repetition).named(field.name)
+ case SQLConf.ParquetOutputTimestampType.TIMESTAMP_MICROS =>
+ Types
+ .primitive(INT64, repetition)
+ .as(LogicalTypeAnnotation.timestampType(true, TimeUnit.MICROS))
+ .named(field.name)
+ case SQLConf.ParquetOutputTimestampType.TIMESTAMP_MILLIS =>
+ Types
+ .primitive(INT64, repetition)
+ .as(LogicalTypeAnnotation.timestampType(true, TimeUnit.MILLIS))
+ .named(field.name)
+ }
+
+ case TimestampNTZType =>
+ Types
+ .primitive(INT64, repetition)
+ .as(LogicalTypeAnnotation.timestampType(false, TimeUnit.MICROS))
+ .named(field.name)
+
+ case BinaryType =>
+ Types.primitive(BINARY, repetition).named(field.name)
+
+ // ======================
+ // Decimals (legacy mode)
+ // ======================
+
+ // Spark 1.4.x and prior versions only support decimals with a maximum precision of 18 and
+ // always store decimals in fixed-length byte arrays. To keep compatibility with these older
+ // versions, here we convert decimals with all precisions to `FIXED_LEN_BYTE_ARRAY` annotated
+ // by `DECIMAL`.
+ case DecimalType.Fixed(precision, scale) if writeLegacyParquetFormat =>
+ Types
+ .primitive(FIXED_LEN_BYTE_ARRAY, repetition)
+ .as(LogicalTypeAnnotation.decimalType(scale, precision))
+ .length(Decimal.minBytesForPrecision(precision))
+ .named(field.name)
+
+ // ========================
+ // Decimals (standard mode)
+ // ========================
+
+ // Uses INT32 for 1 <= precision <= 9
+ case DecimalType.Fixed(precision, scale)
+ if precision <= Decimal.MAX_INT_DIGITS && !writeLegacyParquetFormat =>
+ Types
+ .primitive(INT32, repetition)
+ .as(LogicalTypeAnnotation.decimalType(scale, precision))
+ .named(field.name)
+
+ // Uses INT64 for 1 <= precision <= 18
+ case DecimalType.Fixed(precision, scale)
+ if precision <= Decimal.MAX_LONG_DIGITS && !writeLegacyParquetFormat =>
+ Types
+ .primitive(INT64, repetition)
+ .as(LogicalTypeAnnotation.decimalType(scale, precision))
+ .named(field.name)
+
+ // Uses FIXED_LEN_BYTE_ARRAY for all other precisions
+ case DecimalType.Fixed(precision, scale) if !writeLegacyParquetFormat =>
+ Types
+ .primitive(FIXED_LEN_BYTE_ARRAY, repetition)
+ .as(LogicalTypeAnnotation.decimalType(scale, precision))
+ .length(Decimal.minBytesForPrecision(precision))
+ .named(field.name)
+
+ // ===================================
+ // ArrayType and MapType (legacy mode)
+ // ===================================
+
+ // Spark 1.4.x and prior versions convert `ArrayType` with nullable elements into a 3-level
+ // `LIST` structure. This behavior is somewhat a hybrid of parquet-hive and parquet-avro
+ // (1.6.0rc3): the 3-level structure is similar to parquet-hive while the 3rd level element
+ // field name "array" is borrowed from parquet-avro.
+ case ArrayType(elementType, nullable @ true) if writeLegacyParquetFormat =>
+ // group (LIST) {
+ // optional group bag {
+ // repeated array;
+ // }
+ // }
+
+ // This should not use `listOfElements` here because this new method checks if the
+ // element name is `element` in the `GroupType` and throws an exception if not.
+ // As mentioned above, Spark prior to 1.4.x writes `ArrayType` as `LIST` but with
+ // `array` as its element name as below. Therefore, we build manually
+ // the correct group type here via the builder. (See SPARK-16777)
+ Types
+ .buildGroup(repetition)
+ .as(LogicalTypeAnnotation.listType())
+ .addField(
+ Types
+ .buildGroup(REPEATED)
+ // "array" is the name chosen by parquet-hive (1.7.0 and prior version)
+ .addField(convertField(StructField("array", elementType, nullable)))
+ .named("bag"))
+ .named(field.name)
+
+ // Spark 1.4.x and prior versions convert ArrayType with non-nullable elements into a 2-level
+ // LIST structure. This behavior mimics parquet-avro (1.6.0rc3). Note that this case is
+ // covered by the backwards-compatibility rules implemented in `isElementType()`.
+ case ArrayType(elementType, nullable @ false) if writeLegacyParquetFormat =>
+ // group (LIST) {
+ // repeated element;
+ // }
+
+ // Here too, we should not use `listOfElements`. (See SPARK-16777)
+ Types
+ .buildGroup(repetition)
+ .as(LogicalTypeAnnotation.listType())
+ // "array" is the name chosen by parquet-avro (1.7.0 and prior version)
+ .addField(convertField(StructField("array", elementType, nullable), REPEATED))
+ .named(field.name)
+
+ // Spark 1.4.x and prior versions convert MapType into a 3-level group annotated by
+ // MAP_KEY_VALUE. This is covered by `convertGroupField(field: GroupType): DataType`.
+ case MapType(keyType, valueType, valueContainsNull) if writeLegacyParquetFormat =>
+ // group (MAP) {
+ // repeated group map (MAP_KEY_VALUE) {
+ // required key;
+ // value;
+ // }
+ // }
+ ConversionPatterns.mapType(
+ repetition,
+ field.name,
+ convertField(StructField("key", keyType, nullable = false)),
+ convertField(StructField("value", valueType, valueContainsNull)))
+
+ // =====================================
+ // ArrayType and MapType (standard mode)
+ // =====================================
+
+ case ArrayType(elementType, containsNull) if !writeLegacyParquetFormat =>
+ // group (LIST) {
+ // repeated group list {
+ // element;
+ // }
+ // }
+ Types
+ .buildGroup(repetition)
+ .as(LogicalTypeAnnotation.listType())
+ .addField(
+ Types
+ .repeatedGroup()
+ .addField(convertField(StructField("element", elementType, containsNull)))
+ .named("list"))
+ .named(field.name)
+
+ case MapType(keyType, valueType, valueContainsNull) =>
+ // group (MAP) {
+ // repeated group key_value {
+ // required key;
+ // value;
+ // }
+ // }
+ Types
+ .buildGroup(repetition)
+ .as(LogicalTypeAnnotation.mapType())
+ .addField(
+ Types
+ .repeatedGroup()
+ .addField(convertField(StructField("key", keyType, nullable = false)))
+ .addField(convertField(StructField("value", valueType, valueContainsNull)))
+ .named("key_value"))
+ .named(field.name)
+
+ // ===========
+ // Other types
+ // ===========
+
+ case StructType(fields) =>
+ fields
+ .foldLeft(Types.buildGroup(repetition)) { (builder, field) =>
+ builder.addField(convertField(field))
+ }
+ .named(field.name)
+
+ case udt: UserDefinedType[_] =>
+ convertField(field.copy(dataType = udt.sqlType))
+
+ case _ =>
+ throw QueryCompilationErrors.cannotConvertDataTypeToParquetTypeError(field)
+ }
+ }
+}
diff --git a/common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
new file mode 100644
index 000000000..9e0541f44
--- /dev/null
+++ b/common/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala
@@ -0,0 +1,164 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.spark.sql.comet.util
+
+import java.io.File
+
+import scala.collection.JavaConverters._
+
+import org.apache.arrow.vector.complex.MapVector
+import org.apache.arrow.vector.types._
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
+import org.apache.spark.sql.types._
+
+object Utils {
+ def getConfPath(confFileName: String): String = {
+ sys.env
+ .get("COMET_CONF_DIR")
+ .map { t => new File(s"$t${File.separator}$confFileName") }
+ .filter(_.isFile)
+ .map(_.getAbsolutePath)
+ .orNull
+ }
+
+ def stringToSeq(str: String): Seq[String] = {
+ str.split(",").map(_.trim()).filter(_.nonEmpty)
+ }
+
+ def fromArrowField(field: Field): DataType = {
+ field.getType match {
+ case _: ArrowType.Map =>
+ val elementField = field.getChildren.get(0)
+ val keyType = fromArrowField(elementField.getChildren.get(0))
+ val valueType = fromArrowField(elementField.getChildren.get(1))
+ MapType(keyType, valueType, elementField.getChildren.get(1).isNullable)
+ case ArrowType.List.INSTANCE =>
+ val elementField = field.getChildren().get(0)
+ val elementType = fromArrowField(elementField)
+ ArrayType(elementType, containsNull = elementField.isNullable)
+ case ArrowType.Struct.INSTANCE =>
+ val fields = field.getChildren().asScala.map { child =>
+ val dt = fromArrowField(child)
+ StructField(child.getName, dt, child.isNullable)
+ }
+ StructType(fields.toSeq)
+ case arrowType => fromArrowType(arrowType)
+ }
+ }
+
+ def fromArrowType(dt: ArrowType): DataType = dt match {
+ case ArrowType.Bool.INSTANCE => BooleanType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 => ByteType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 2 => ShortType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 4 => IntegerType
+ case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 8 => LongType
+ case float: ArrowType.FloatingPoint if float.getPrecision == FloatingPointPrecision.SINGLE =>
+ FloatType
+ case float: ArrowType.FloatingPoint if float.getPrecision == FloatingPointPrecision.DOUBLE =>
+ DoubleType
+ case ArrowType.Utf8.INSTANCE => StringType
+ case ArrowType.Binary.INSTANCE => BinaryType
+ case _: ArrowType.FixedSizeBinary => BinaryType
+ case d: ArrowType.Decimal => DecimalType(d.getPrecision, d.getScale)
+ case date: ArrowType.Date if date.getUnit == DateUnit.DAY => DateType
+ case ts: ArrowType.Timestamp if ts.getUnit == TimeUnit.MICROSECOND => TimestampType
+ case ArrowType.Null.INSTANCE => NullType
+ case yi: ArrowType.Interval if yi.getUnit == IntervalUnit.YEAR_MONTH =>
+ YearMonthIntervalType()
+ case di: ArrowType.Interval if di.getUnit == IntervalUnit.DAY_TIME => DayTimeIntervalType()
+ case _ => throw new UnsupportedOperationException(s"Unsupported data type: ${dt.toString}")
+ }
+
+ /** Maps data type from Spark to Arrow. NOTE: timeZoneId required for TimestampTypes */
+ def toArrowType(dt: DataType, timeZoneId: String): ArrowType =
+ dt match {
+ case BooleanType => ArrowType.Bool.INSTANCE
+ case ByteType => new ArrowType.Int(8, true)
+ case ShortType => new ArrowType.Int(8 * 2, true)
+ case IntegerType => new ArrowType.Int(8 * 4, true)
+ case LongType => new ArrowType.Int(8 * 8, true)
+ case FloatType => new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)
+ case DoubleType => new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)
+ case StringType => ArrowType.Utf8.INSTANCE
+ case BinaryType => ArrowType.Binary.INSTANCE
+ case DecimalType.Fixed(precision, scale) => new ArrowType.Decimal(precision, scale, 128)
+ case DateType => new ArrowType.Date(DateUnit.DAY)
+ case TimestampType =>
+ if (timeZoneId == null) {
+ throw new UnsupportedOperationException(
+ s"${TimestampType.catalogString} must supply timeZoneId parameter")
+ } else {
+ new ArrowType.Timestamp(TimeUnit.MICROSECOND, timeZoneId)
+ }
+ case TimestampNTZType =>
+ new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)
+ case _ =>
+ throw new UnsupportedOperationException(s"Unsupported data type: ${dt.catalogString}")
+ }
+
+ /** Maps field from Spark to Arrow. NOTE: timeZoneId required for TimestampType */
+ def toArrowField(name: String, dt: DataType, nullable: Boolean, timeZoneId: String): Field = {
+ dt match {
+ case ArrayType(elementType, containsNull) =>
+ val fieldType = new FieldType(nullable, ArrowType.List.INSTANCE, null)
+ new Field(
+ name,
+ fieldType,
+ Seq(toArrowField("element", elementType, containsNull, timeZoneId)).asJava)
+ case StructType(fields) =>
+ val fieldType = new FieldType(nullable, ArrowType.Struct.INSTANCE, null)
+ new Field(
+ name,
+ fieldType,
+ fields
+ .map { field =>
+ toArrowField(field.name, field.dataType, field.nullable, timeZoneId)
+ }
+ .toSeq
+ .asJava)
+ case MapType(keyType, valueType, valueContainsNull) =>
+ val mapType = new FieldType(nullable, new ArrowType.Map(false), null)
+ // Note: Map Type struct can not be null, Struct Type key field can not be null
+ new Field(
+ name,
+ mapType,
+ Seq(
+ toArrowField(
+ MapVector.DATA_VECTOR_NAME,
+ new StructType()
+ .add(MapVector.KEY_NAME, keyType, nullable = false)
+ .add(MapVector.VALUE_NAME, valueType, nullable = valueContainsNull),
+ nullable = false,
+ timeZoneId)).asJava)
+ case dataType =>
+ val fieldType = new FieldType(nullable, toArrowType(dataType, timeZoneId), null)
+ new Field(name, fieldType, Seq.empty[Field].asJava)
+ }
+ }
+
+ /**
+ * Maps schema from Spark to Arrow. NOTE: timeZoneId required for TimestampType in StructType
+ */
+ def toArrowSchema(schema: StructType, timeZoneId: String): Schema = {
+ new Schema(schema.map { field =>
+ toArrowField(field.name, field.dataType, field.nullable, timeZoneId)
+ }.asJava)
+ }
+}
diff --git a/common/src/test/java/org/apache/comet/parquet/TestColumnReader.java b/common/src/test/java/org/apache/comet/parquet/TestColumnReader.java
new file mode 100644
index 000000000..d4e748a9b
--- /dev/null
+++ b/common/src/test/java/org/apache/comet/parquet/TestColumnReader.java
@@ -0,0 +1,193 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.comet.parquet;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.BiFunction;
+
+import scala.collection.JavaConverters;
+
+import org.junit.Test;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FixedSizeBinaryVector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
+import org.apache.spark.sql.types.*;
+import org.apache.spark.sql.vectorized.ColumnVector;
+
+import org.apache.comet.vector.CometPlainVector;
+import org.apache.comet.vector.CometVector;
+
+import static org.apache.spark.sql.types.DataTypes.*;
+import static org.junit.Assert.*;
+
+@SuppressWarnings("unchecked")
+public class TestColumnReader {
+ private static final int BATCH_SIZE = 1024;
+ private static final List TYPES =
+ Arrays.asList(
+ BooleanType,
+ ByteType,
+ ShortType,
+ IntegerType,
+ LongType,
+ FloatType,
+ DoubleType,
+ BinaryType,
+ DecimalType.apply(5, 2),
+ DecimalType.apply(18, 10),
+ DecimalType.apply(19, 5));
+ private static final List
