Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Spark/NativeIO] Configure LSH compute by columnar metadata #571

Merged
merged 1 commit into from
Jan 14, 2025
Merged

[Spark/NativeIO] Configure LSH compute by columnar metadata #571

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions lakesoul-common/src/main/scala/com/dmetasoul/lakesoul/meta/LakeSoulOptions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@ object LakeSoulOptions {
val TIME_ZONE = "timezone"
val DISCOVERY_INTERVAL = "discoveryinterval"

object SchemaFieldMetadata {
val LSH_EMBEDDING_DIMENSION = "lsh_embedding_dimension"
val LSH_BIT_WIDTH = "lsh_bit_width"
val LSH_RNG_SEED = "lsh_rng_seed"
}

object ReadType extends Enumeration {
val FULL_READ = "fullread"
val SNAPSHOT_READ = "snapshot"
Expand Down
412 changes: 210 additions & 202 deletions lakesoul-spark/src/test/scala/org/apache/spark/sql/lakesoul/commands/MergeIntoSQLSuite.scala
View file
Open in desktop

Large diffs are not rendered by default.

21 changes: 16 additions & 5 deletions native-io/lakesoul-io-java/src/main/scala/org/apache/spark/sql/arrow/ArrowUtils.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

package org.apache.spark.sql.arrow

import com.dmetasoul.lakesoul.meta.LakeSoulOptions
import com.dmetasoul.lakesoul.meta.LakeSoulOptions.SchemaFieldMetadata.{LSH_BIT_WIDTH, LSH_EMBEDDING_DIMENSION, LSH_RNG_SEED}
import org.apache.arrow.memory.RootAllocator
import org.apache.arrow.vector.complex.MapVector
import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
Expand Down Expand Up @@ -167,11 +169,20 @@ object ArrowUtils {
def toArrowSchema(schema: StructType, timeZoneId: String = "UTC"): Schema = {
new Schema(schema.map { field =>
val comment = field.getComment
val metadata = if (comment.isDefined) {
val map = new util.HashMap[String, String]
map.put("spark_comment", comment.get)
map
} else null
val metadata = new util.HashMap[String, String]
if (field.metadata.contains(LSH_EMBEDDING_DIMENSION)) {
metadata.put(LSH_EMBEDDING_DIMENSION, field.metadata.getString(LSH_EMBEDDING_DIMENSION))
}
if (field.metadata.contains(LSH_BIT_WIDTH)) {
metadata.put(LSH_BIT_WIDTH, field.metadata.getString(LSH_BIT_WIDTH))
}
if (field.metadata.contains(LSH_RNG_SEED)) {
metadata.put(LSH_RNG_SEED, field.metadata.getString(LSH_RNG_SEED))
}

if (comment.isDefined) {
metadata.put("spark_comment", comment.get)
}
toArrowField(field.name, field.dataType, field.nullable, timeZoneId, field.metadata, metadata)
}.asJava)
}
Expand Down
16 changes: 4 additions & 12 deletions rust/lakesoul-io/src/lakesoul_io_config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,8 @@ pub static OPTION_KEY_HASH_BUCKET_ID: &str = "hash_bucket_id";
pub static OPTION_KEY_CDC_COLUMN: &str = "cdc_column";
pub static OPTION_KEY_IS_COMPACTED: &str = "is_compacted";
pub static OPTION_KEY_MAX_FILE_SIZE: &str = "max_file_size";
pub static OPTION_KEY_IS_LSH: &str = "is_lsh";
pub static OPTION_KEY_NBITS: &str = "nbits";
pub static OPTION_KEY_D: &str= "d";
pub static OPTION_KEY_COMPUTE_LSH: &str = "compute_lsh";


#[derive(Debug, Derivative)]
#[derivative(Default, Clone)]
Expand Down Expand Up @@ -196,17 +195,10 @@ impl LakeSoulIOConfig {
self.option(OPTION_KEY_IS_COMPACTED).map_or(false, |x| x.eq("true"))
}

pub fn is_lsh(&self) -> bool {
self.option(OPTION_KEY_IS_LSH).map_or(false,|x| x.eq("true"))
}

pub fn nbits(&self) -> Option<u64>{
self.option(OPTION_KEY_NBITS).map(|x| x.parse().unwrap())
pub fn compute_lsh(&self) -> bool {
self.option(OPTION_KEY_COMPUTE_LSH).map_or(true,|x| x.eq("true"))
}

pub fn d(&self) -> Option<u64>{
self.option(OPTION_KEY_D).map(|x| x.parse().unwrap())
}
}

#[derive(Derivative, Debug)]
Expand Down
Loading
Loading