kazuyukitanimura commented on code in PR #1047:
URL: https://github.com/apache/datafusion-comet/pull/1047#discussion_r1826350543
##########
common/src/main/java/org/apache/comet/parquet/ColumnReader.java:
##########
@@ -201,53 +207,56 @@ public CometDecodedVector loadVector() {
boolean isUuid =
logicalTypeAnnotation instanceof
LogicalTypeAnnotation.UUIDLogicalTypeAnnotation;
- long[] addresses = Native.currentBatch(nativeHandle);
+ array = ArrowArray.allocateNew(ALLOCATOR);
+ schema = ArrowSchema.allocateNew(ALLOCATOR);
- try (ArrowArray array = ArrowArray.wrap(addresses[0]);
- ArrowSchema schema = ArrowSchema.wrap(addresses[1])) {
- FieldVector vector = importer.importVector(array, schema);
+ long arrayAddr = array.memoryAddress();
+ long schemaAddr = schema.memoryAddress();
- DictionaryEncoding dictionaryEncoding =
vector.getField().getDictionary();
+ Native.currentBatch(nativeHandle, arrayAddr, schemaAddr);
- CometPlainVector cometVector = new CometPlainVector(vector,
useDecimal128);
+ FieldVector vector = importer.importVector(array, schema);
- // 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();
+ DictionaryEncoding dictionaryEncoding = vector.getField().getDictionary();
- 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;
- }
+ 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();
- // We should already re-initiate `CometDictionary` here because
`Data.importVector` API will
- // release the previous dictionary vector and create a new one.
- Dictionary arrowDictionary =
importer.getProvider().lookup(dictionaryEncoding.getId());
- CometPlainVector dictionaryVector =
- new CometPlainVector(arrowDictionary.getVector(), useDecimal128,
isUuid);
+ if (dictionaryEncoding == null) {
if (dictionary != null) {
- dictionary.setDictionaryVector(dictionaryVector);
- } else {
- dictionary = new CometDictionary(dictionaryVector);
+ // 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;
}
-
- currentVector =
- new CometDictionaryVector(
- cometVector, dictionary, importer.getProvider(), useDecimal128,
false, isUuid);
-
- currentVector =
- new CometDictionaryVector(cometVector, dictionary,
importer.getProvider(), useDecimal128);
+ // 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;
}
+
+ // We should already re-initiate `CometDictionary` here because
`Data.importVector` API will
+ // release the previous dictionary vector and create a new one.
+ Dictionary arrowDictionary =
importer.getProvider().lookup(dictionaryEncoding.getId());
+ CometPlainVector dictionaryVector =
+ new CometPlainVector(arrowDictionary.getVector(), useDecimal128,
isUuid);
+ if (dictionary != null) {
+ dictionary.setDictionaryVector(dictionaryVector);
+ } else {
+ dictionary = new CometDictionary(dictionaryVector);
+ }
+
+ currentVector =
+ new CometDictionaryVector(
+ cometVector, dictionary, importer.getProvider(), useDecimal128,
false, isUuid);
+
+ currentVector =
+ new CometDictionaryVector(cometVector, dictionary,
importer.getProvider(), useDecimal128);
Review Comment:
Unrelated but we should not be needing to call `CometDictionaryVector` twice
##########
native/core/src/parquet/mod.rs:
##########
@@ -539,24 +537,16 @@ pub extern "system" fn
Java_org_apache_comet_parquet_Native_currentBatch(
e: JNIEnv,
_jclass: JClass,
handle: jlong,
-) -> jlongArray {
- try_unwrap_or_throw(&e, |env| {
+ array_addr: jlong,
+ schema_addr: jlong,
+) {
+ try_unwrap_or_throw(&e, |_env| {
let ctx = get_context(handle)?;
let reader = &mut ctx.column_reader;
let data = reader.current_batch();
- let (array, schema) = data.to_spark()?;
+ data.move_to_spark(array_addr, schema_addr)?;
Review Comment:
nit: maybe we can just return `data.move_to_spark(array_addr, schema_addr)`
so that we do not need `ok(())`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
