Skip to content

Commit

Permalink
feat: support for vector; quick mode for geo && vector data for dump
Browse files Browse the repository at this point in the history
  • Loading branch information
yongcai committed Dec 27, 2024
1 parent 43f0044 commit 8568966
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 7 deletions.
2 changes: 2 additions & 0 deletions src/pyinnodb/cli/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,8 @@ def transfter(nd):
elif isinstance(field, MGeo):
d = field.build().hex() # .zfill(50)
vs.append("0x" + d)
elif isinstance(field, bytes):
vs.append("0x"+field.hex())
else:
vs.append(repr(field))
values.append(f"({','.join(vs)})")
Expand Down
5 changes: 5 additions & 0 deletions src/pyinnodb/const/dd_column_type.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ class DDColumnType(Enum):
STRING = 29
GEOMETRY = 30
JSON = 31
VECTOR = 32

def is_int_number(self):
return self in _int_number_type
Expand Down Expand Up @@ -88,13 +89,16 @@ def is_big(cls, t):
DDColumnType.JSON,
DDColumnType.TINY_BLOB,
DDColumnType.GEOMETRY,
DDColumnType.VECTOR,
]

_big_type = [
DDColumnType.MEDIUM_BLOB,
DDColumnType.LONG_BLOB,
DDColumnType.BLOB,
DDColumnType.JSON,
DDColumnType.GEOMETRY,
DDColumnType.VECTOR,
]

DDColConf = namedtuple("DDColConf", "type size")
Expand Down Expand Up @@ -132,6 +136,7 @@ class DDColConf(DDColConf, Enum):
STRING = DDColumnType.STRING, 0
GEOMETRY = DDColumnType.GEOMETRY, 0
JSON = DDColumnType.JSON, 0
VECTOR = DDColumnType.VECTOR, 0

@classmethod
def get_col_type_conf(cls, type) -> DDColConf:
Expand Down
4 changes: 2 additions & 2 deletions src/pyinnodb/disk_struct/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ class MIndexPage(CC):
system_records: MIndexSystemRecord = cfield(MIndexSystemRecord)

@classmethod
def default_value_parser(cls, dd_object: Table, transfter=None, hidden_col=False):
def default_value_parser(cls, dd_object: Table, transfter=None, hidden_col=False, quick=True):
primary_data_layout_col = dd_object.get_disk_data_layout()

def value_parser(rh: MRecordHeader, f):
Expand Down Expand Up @@ -227,7 +227,7 @@ def value_parser(rh: MRecordHeader, f):
if col.ordinal_position in null_col_data:
col_value = None
else:
col_value = col.read_data(f, vs)
col_value = col.read_data(f, vs, quick=quick)
except Exception as e:
print("cur before is ", cur_before, vs, col)
raise e
Expand Down
23 changes: 18 additions & 5 deletions src/pyinnodb/sdi/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,7 +395,7 @@ def _read_varchar(self, stream, size):
else:
return data

def read_data(self, stream, size=None):
def read_data(self, stream, size=None, quick=True):
if self.name == "DB_ROLL_PTR":
return MRollbackPointer.parse_stream(stream)
dtype = DDColumnType(self.type)
Expand All @@ -405,6 +405,15 @@ def read_data(self, stream, size=None):
dsize = self.size
if dtype.is_int_number():
return self._read_int(stream, dsize)
elif dtype == DDColumnType.VECTOR:
if quick:
return stream.read(dsize)
else:
vec = []
for i in range(int(dsize / 4)):
byte_data = stream.read(4)
vec.append(struct.unpack("f", byte_data)[0])
return vec
elif dtype == DDColumnType.FLOAT:
byte_data = stream.read(dsize)
if dsize == 4:
Expand Down Expand Up @@ -476,9 +485,12 @@ def read_data(self, stream, size=None):
except Exception as e:
return data
elif dtype == DDColumnType.GEOMETRY:
data = MGeo.parse_stream(stream)
logging.debug("geometry data is %s, size is %d", data, dsize)
return data
if quick:
return stream.read(dsize)
else:
data = MGeo.parse_stream(stream)
logging.debug("geometry data is %s, size is %d", data, dsize)
return data


decimal_leftover_part = {
Expand Down Expand Up @@ -846,7 +858,8 @@ def search(self, f, primary_key, hidden_col):
else:
primary_key = self.build_primary_key_bytes((primary_key,))
value_parser = MIndexPage.default_value_parser(
self, hidden_col=hidden_col, transfter=lambda id: id
self, hidden_col=hidden_col, transfter=lambda id: id,
quick=False,
)

while first_leaf_page != 4294967295:
Expand Down

0 comments on commit 8568966

Please sign in to comment.