code: # @Copyright: CEA-LIST/DIASI/SIALV/LVA (2023)
  # @Author: CEA-LIST/DIASI/SIALV/LVA <pixano@cea.fr>
  # @License: CECILL-C
  #
  # This software is a collaborative computer program whose purpose is to
  # generate and explore labeled data for computer vision applications.
  # This software is governed by the CeCILL-C license under French law and
  # abiding by the rules of distribution of free software. You can use,
  # modify and/ or redistribute the software under the terms of the CeCILL-C
  # license as circulated by CEA, CNRS and INRIA at the following URL
  #
  # http://www.cecill.info
  
  from collections import defaultdict
  from pathlib import Path
  from typing import Optional
  
  import duckdb
  import lancedb
  import pyarrow as pa
  import pyarrow.dataset as pa_ds
  from pydantic import BaseModel, ConfigDict
  from s3path import S3Path
  
  from pixano.core import Image
  from pixano.data.dataset.dataset_info import DatasetInfo
  from pixano.data.dataset.dataset_item import DatasetItem
  from pixano.data.dataset.dataset_stat import DatasetStat
  from pixano.data.dataset.dataset_table import DatasetTable
  from pixano.data.fields import Fields, field_to_pyarrow
  from pixano.data.item.item_feature import FeaturesValues, FeatureValues, ItemFeature
  from pixano.data.item.item_object import ItemObject
  
  
  class Dataset(BaseModel):
      """Dataset
  
      Attributes:
          path (Path | S3Path): Dataset path
          info (DatasetInfo, optional): Dataset info
          stats (list[DatasetStat], optional): Dataset stats
          thumbnail (str, optional): Dataset thumbnail base 64 URL
      """
  
      path: Path | S3Path
      info: Optional[DatasetInfo] = None
      stats: Optional[list[DatasetStat]] = None
      thumbnail: Optional[str] = None
      # Allow arbitrary types because of S3 Path
      model_config = ConfigDict(arbitrary_types_allowed=True)
  
      def __init__(
          self,
          path: Path | S3Path,
      ):
          """Initialize dataset
  
          Args:
              path (Path | S3Path): Dataset path
          """
  
          info_file = path / "db.json"
          stats_file = path / "stats.json"
          thumb_file = path / "preview.png"
  
          # Define public attributes through Pydantic BaseModel
          super().__init__(
              path=path,
              info=DatasetInfo.from_json(info_file),
              stats=DatasetStat.from_json(stats_file) if stats_file.is_file() else None,
              thumbnail=(
                  Image(uri=thumb_file.absolute().as_uri()).url
                  if thumb_file.is_file()
                  else None
              ),
          )
  
      @property
      def media_dir(self) -> Path | S3Path:
          """Return dataset media directory
  
          Returns:
              Path | S3Path: Dataset media directory
          """
  
          return self.path / "media"
  
      @property
      def num_rows(self) -> int:
          """Return number of rows in dataset
  
          Returns:
              int: Number of rows
          """
  
          ds = self.connect()
  
          # Return number of rows of main table
          return len(ds.open_table("db"))
  
      def get_item_uuid(self, original_id: str) -> str:
          """Get id (uuid) from original data item id, if exist
  
          Args:
              original_id (str): original data item id
  
          Returns:
              str: unique item id, or original_id if not found
          """
          ds_tables = self.open_tables()
          # pylint: disable=unused-variable
          lance_table = ds_tables["main"]["db"].to_lance()
          result = (
              duckdb.query(
                  f"SELECT id FROM lance_table WHERE original_id = '{original_id}'"
              )
              .to_arrow_table()
              .to_pylist()
          )
          if result and len(result) == 1:
              return result[0]["id"]
          return original_id
  
      def get_object_uuid(self, original_id: str) -> str:
          """Get id (uuid) from original data object id, if exist
  
          Args:
              original_id (str): original data object id
  
          Returns:
              str: unique item id, or original_id if not found
          """
          ds_tables = self.open_tables()
          # pylint: disable=unused-variable
          lance_table = ds_tables["objects"]["objects"].to_lance()
          result = (
              duckdb.query(
                  f"SELECT id FROM lance_table WHERE original_id = '{original_id}'"
              )
              .to_arrow_table()
              .to_pylist()
          )
          if result and len(result) == 1:
              return result[0]["id"]
          return original_id
  
      def load_info(
          self,
          load_stats: bool = False,
          load_thumbnail: bool = False,
          load_features_values: bool = False,
      ) -> DatasetInfo:
          """Return dataset info with thumbnail and stats inside
  
          Args:
              load_stats (bool, optional): Load dataset stats. Defaults to False.
              load_thumbnail (bool, optional): Load dataset thumbnail. Defaults to False.
              load_features_values (bool, optional): Load available values. Defaults to False.
  
          Returns:
              DatasetInfo: Dataset info
          """
  
          info = DatasetInfo.from_json(
              self.path / "db.json",
              load_stats=load_stats,
              load_thumbnail=load_thumbnail,
          )
          if load_features_values:
              info.features_values = self.get_features_values(info.features_values)
  
          return info
  
      def save_info(self):
          """Save updated dataset info"""
  
          self.info.save(self.path)
  
      @staticmethod
      def create(path: Path | S3Path, info: DatasetInfo) -> "Dataset":
          """Create dataset
  
          Args:
              path (Path | S3Path): Path to create dataset in
              info (DatasetInfo): Dataset info
  
          Returns:
              Dataset: Created dataset
          """
  
          # Create DatasetInfo file
          path.mkdir(parents=True, exist_ok=True)
          info.save(path)
  
          # Load dataset
          dataset = Dataset(path)
  
          # Create dataset tables
          for group_name, table_group in dataset.info.tables.items():
              for table in table_group:
                  dataset.create_table(table, group_name, add_to_info=False)
  
          return dataset
  
      def connect(self) -> lancedb.db.DBConnection:
          """Connect to dataset with LanceDB
  
          Returns:
              lancedb.db.DBConnection: Dataset LanceDB connection
          """
  
          if isinstance(self.path, S3Path):
              return lancedb.connect(self.path.as_uri())
          return lancedb.connect(self.path)
  
      def open_tables(self) -> dict[str, dict[str, lancedb.db.LanceTable]]:
          """Open dataset tables with LanceDB
  
          Returns:
              dict[str, dict[str, lancedb.db.LanceTable]]: Dataset tables
          """
  
          ds = self.connect()
  
          ds_tables: dict[str, dict[str, lancedb.db.LanceTable]] = defaultdict(dict)
  
          for group_name, table_group in self.info.tables.items():
              for table in table_group:
                  try:
                      ds_tables[group_name][table.name] = ds.open_table(table.name)
                  except FileNotFoundError as e:
                      # If optional table, remove from DatasetInfo
                      if group_name in ["objects", "embeddings", "active_learning"]:
                          self.info.tables[group_name].remove(table)
                          self.save_info()
                      else:
                          raise FileNotFoundError from e
  
          return ds_tables
  
      def create_table(
          self,
          table: DatasetTable,
          table_group: str,
          add_to_info: bool = True,
      ):
          """Create a new table in the dataset
  
          Args:
              table (DatasetTable): Table to create
              table_group (str): Table group
              add_to_info (bool, optional): Add table to DatasetInfo. Defaults to True.
          """
  
          # Create Lance table
          ds = self.connect()
          # pylint: disable=unexpected-keyword-arg
          ds.create_table(
              table.name,
              schema=Fields(table.fields).to_schema(),
              mode="overwrite",
          )
  
          # Save table to DatasetInfo
          if add_to_info:
              if table_group in self.info.tables:
                  self.info.tables[table_group].append(table)
              else:
                  self.info.tables[table_group] = [table]
              self.save_info()
  
      def update_table(
          self,
          element: DatasetItem | ItemObject,
          table: lancedb.db.LanceTable,
          table_group: str,
          table_name: str,
      ):
          """Update a table with new features or base fields
  
          Args:
              element (DatasetItem | ItemObject): Table element (item or object)
              table (lancedb.db.LanceTable): Table to update
              table_group (str): Table group
              table_name (str): Table name
          """
  
          new_columns: list[ItemFeature] = []
  
          # Check for new base fields
          base_fields = {"review_state": "str", "bbox": "bbox", "mask": "compressedrle"}
          new_columns.extend(
              [
                  ItemFeature(name=field_name, dtype=field_type, value=None)
                  for field_name, field_type in base_fields.items()
                  if hasattr(element, field_name) and field_name not in table.schema.names
              ]
          )
          # Check for new features
          if element.features is not None:
              new_columns.extend(
                  [
                      feat
                      for feat in element.features.values()
                      if feat.name not in table.schema.names
                  ]
              )
          # Add new columns
          if len(new_columns) > 0:
              new_columns_table = table.to_lance().to_table(columns=["id"])
              for col in new_columns:
                  # None should be suported for booleans with pylance 0.9.1
                  # None is not supported for integers and floats yet
                  none_value = (
                      False
                      if col.dtype == "bool"
                      else 0 if col.dtype in ("int", "float") else None
                  )
                  col_array = pa.array(
                      [none_value] * len(table),
                      type=field_to_pyarrow(col.dtype),
                  )
                  new_columns_table = new_columns_table.append_column(
                      pa.field(col.name, field_to_pyarrow(col.dtype)), col_array
                  )
                  # Update DatasetInfo
                  for info_table in self.info.tables[table_group]:
                      if info_table.name == table_name:
                          info_table.fields[col.name] = col.dtype
  
              # Merge with main table
              table.to_lance().merge(new_columns_table, "id")
              self.save_info()
  
      def load_items(
          self,
          limit: int,
          offset: int,
          load_active_learning: bool = True,
      ) -> list[DatasetItem]:
          """Load dataset items in selected tables
  
          Args:
              limit (int): Items limit
              offset (int): Items offset
              load_active_learning (bool, optional): Load items active learning info. Defaults to True.
          Returns:
              list[DatasetItem]: List of dataset items
          """
  
          # Update info in case of change
          self.info = self.load_info()
  
          # Load tables
          ds_tables = self.open_tables()
  
          # Load PyArrow items from tables
          pyarrow_items: dict[str, dict[str, pa.Table]] = defaultdict(dict)
  
          # Load PyArrow items from main table
          # pylint: disable=unused-variable
          lance_table = ds_tables["main"]["db"].to_lance()
          id_field = "original_id" if "original_id" in lance_table.schema.names else "id"
          pyarrow_items["main"]["db"] = duckdb.query(
              f"SELECT * FROM lance_table ORDER BY len({id_field}), {id_field} LIMIT {limit} OFFSET {offset}"
          ).to_arrow_table()
          id_list = tuple(pyarrow_items["main"]["db"].to_pydict()["id"])
  
          # Media tables
          for media_source, media_table in ds_tables["media"].items():
              lance_table = media_table.to_lance()
              pyarrow_items["media"][media_source] = duckdb.query(
                  f"SELECT * FROM lance_table WHERE id IN {id_list}"
              ).to_arrow_table()
  
          # Active Learning tables
          if load_active_learning:
              for al_source, al_table in ds_tables["active_learning"].items():
                  lance_table = al_table.to_lance()
                  pyarrow_items["active_learning"][al_source] = duckdb.query(
                      f"SELECT * FROM lance_table  WHERE id IN {id_list}"
                  ).to_arrow_table()
  
          if pyarrow_items["main"]["db"].num_rows > 0:
              # Split results
              pyarrow_item_list = self._split_items(pyarrow_items, load_active_learning)
  
              return [
                  DatasetItem.from_pyarrow(pyarrow_item, self.info, self.media_dir)
                  for pyarrow_item in pyarrow_item_list
              ]
          return None
  
      def search_items(
          self,
          limit: int,
          offset: int,
          query: dict[str, str],
          load_active_learning: bool = True,
      ):
          """Search for dataset items in selected tables
  
          Args:
              limit (int): Items limit
              offset (int): Items offset
              query (dict[str, str]): Search query
              load_active_learning (bool, optional): Load items active learning info. Defaults to True.
          Returns:
              list[DatasetItem]: List of dataset items
          """
  
          # Update info in case of change
          self.info = self.load_info()
  
          # Load PyArrow items from tables
          pyarrow_items: dict[str, dict[str, pa.Table]] = defaultdict(dict)
  
          # Search items with selected method
          if query["model"] in ["CLIP"]:
              pyarrow_items = self._embeddings_search(limit, offset, query)
          # NOTE: metadata search could go here
  
          if pyarrow_items is not None and pyarrow_items["main"]["db"].num_rows > 0:
              # Split results
              pyarrow_item_list = self._split_items(pyarrow_items, load_active_learning)
  
              return [
                  DatasetItem.from_pyarrow(pyarrow_item, self.info, self.media_dir)
                  for pyarrow_item in pyarrow_item_list
              ]
          return None
  
      def _embeddings_search(
          self,
          limit: int,
          offset: int,
          query: dict[str, str],
      ) -> dict[str, dict[str, pa.Table]]:
          """Perform item semantic search with embeddings
  
          Args:
              limit (int): Items limit
              offset (int): Items offset
              query (dict[str, str]): Search query
  
          Raises:
              ImportError: Required pixano-inference module could not be imported
  
          Returns:
              dict[str, dict[str, pa.Table]]: Search results
          """
  
          # Load tables
          ds_tables = self.open_tables()
  
          # Create PyArrow items
          pyarrow_items: dict[str, dict[str, pa.Table]] = defaultdict(dict)
  
          # Find CLIP embeddings
          if "embeddings" not in self.info.tables:
              return None
          for table in self.info.tables["embeddings"]:
              if table.type == "search" and table.source == query["model"]:
                  sem_search_table = ds_tables["embeddings"][table.name]
                  sem_search_views = [
                      field_name
                      for field_name, field_type in table.fields.items()
                      if field_type == "vector(512)"
                  ]
  
          if query["model"] == "CLIP":
              # Initialize CLIP model
              try:
                  # pylint: disable=import-outside-toplevel
                  from pixano_inference.transformers import CLIP
              except ImportError as e:
                  raise ImportError(
                      "Please install the pixano-inference module to perform semantic search with CLIP"
                  ) from e
  
              model = CLIP()
  
          model_query = model.semantic_search(query["search"])
  
          # Perform semantic search
          # pylint: disable=unused-variable
          results_table = (
              sem_search_table.search(model_query, sem_search_views[0])
              .limit(min(offset + limit, self.num_rows))
              .to_arrow()
          )
  
          # If more than one view, search on all views and select the best results based on distance
          if len(sem_search_views) > 1:
              for view in sem_search_views[1:]:
                  view_results_table = (
                      sem_search_table.search(model_query, view)
                      .limit(min(offset + limit, self.num_rows))
                      .to_arrow()
                  )
                  results_table = duckdb.query(
                      "SELECT id, results_table._distance as distance_1, view_results_table._distance as distance_2 FROM results_table LEFT JOIN view_results_table USING (id)"
                  ).to_arrow_table()
  
                  results_table = duckdb.query(
                      "SELECT (id), (SELECT Min(v) FROM (VALUES (distance_1), (distance_2)) AS value(v)) as _distance FROM results_table"
                  ).to_arrow_table()
  
          # Filter results to page
          results_table = duckdb.query(
              f"SELECT id, _distance as distance FROM results_table ORDER BY _distance ASC LIMIT {limit} OFFSET {offset}"
          ).to_arrow_table()
  
          # Join with main table
          main_table = ds_tables["main"]["db"].to_lance()
          pyarrow_items["main"]["db"] = duckdb.query(
              "SELECT * FROM results_table LEFT JOIN main_table USING (id) ORDER BY distance ASC"
          ).to_arrow_table()
  
          return pyarrow_items
  
      def _split_items(
          self,
          pyarrow_items: dict[str, dict[str, pa.Table]],
          load_active_learning: bool,
      ) -> list[dict[str, dict[str, pa.Table]]]:
          """Split PyArrow tables into list of PyArrow tables
  
          Args:
              pyarrow_items (dict[str, dict[str, pa.Table]]): PyArrow tables
              load_active_learning (bool): Load items active learning info
  
          Returns:
              list[dict[str, dict[str, pa.Table]]]: List of PyArrow tables
          """
  
          # Load tables
          ds_tables = self.open_tables()
  
          # Create list of PyArrow tables
          pyarrow_item_list: list[dict[str, dict[str, pa.Table]]] = []
  
          for index in range(pyarrow_items["main"]["db"].num_rows):
              pyarrow_item_list.append(defaultdict(dict))
              # Main table
              pyarrow_item_list[index]["main"]["db"] = pyarrow_items["main"]["db"].take(
                  [index]
              )
              item_id = pyarrow_item_list[index]["main"]["db"].to_pylist()[0]["id"]
  
              # Media tables
              for media_source, media_table in ds_tables["media"].items():
                  # If media table already created
                  if "media" in pyarrow_items:
                      pyarrow_item_list[index]["media"][media_source] = (
                          pa_ds.dataset(pyarrow_items["media"][media_source])
                          .scanner(filter=pa_ds.field("id") == item_id)
                          .to_table()
                      )
                  # Else, retrieve media items individually
                  else:
                      lance_scanner = media_table.to_lance().scanner(
                          filter=f"id in ('{item_id}')"
                      )
                      pyarrow_item_list[index]["media"][
                          media_source
                      ] = lance_scanner.to_table()
  
              # Active learning tables
              if load_active_learning:
                  # If active learning table already created
                  if "active_learning" in pyarrow_items:
                      for al_source in ds_tables["active_learning"].keys():
                          pyarrow_item_list[index]["active_learning"][al_source] = (
                              pa_ds.dataset(pyarrow_items["active_learning"][al_source])
                              .scanner(filter=pa_ds.field("id") == item_id)
                              .to_table()
                          )
                  # Else, retrieve active learning items individually
                  else:
                      for al_source, al_table in ds_tables["active_learning"].items():
                          lance_scanner = al_table.to_lance().scanner(
                              filter=f"id in ('{item_id}')"
                          )
                          pyarrow_item_list[index]["active_learning"][
                              al_source
                          ] = lance_scanner.to_table()
  
          return pyarrow_item_list
  
      def load_item(
          self,
          item_id: str,
          load_media: bool = True,
          load_objects: bool = False,
          load_active_learning: bool = True,
          load_embeddings: bool = False,
          model_id: str = None,
      ) -> DatasetItem:
          """Find dataset item in selected tables
  
          Args:
              item_id (str): Dataset item ID
              load_media (bool, optional): Load item media. Defaults to True.
              load_objects (bool, optional): Load item objects. Defaults to False.
              load_active_learning (bool, optional): Load item active learning info. Defaults to True.
              load_embeddings (bool, optional): Load item embeddings. Defaults to False.
              model_id (str, optional): Model ID (ONNX file path) of embeddings to load. Defaults to None.
          Returns:
              DatasetItem: Dataset item
          """
  
          # Update info in case of change
          self.info = self.load_info()
  
          # Load tables
          ds_tables = self.open_tables()
  
          # Load PyArrow item from tables
          pyarrow_item: dict[str, dict[str, pa.Table]] = defaultdict(dict)
  
          # Load PyArrow item from main table
          lance_scanner = (
              ds_tables["main"]["db"].to_lance().scanner(filter=f"id in ('{item_id}')")
          )
          pyarrow_item["main"]["db"] = lance_scanner.to_table()
  
          # Load PyArrow item from media tables
          if load_media:
              for table_name, media_table in ds_tables["media"].items():
                  lance_scanner = media_table.to_lance().scanner(
                      filter=f"id in ('{item_id}')"
                  )
                  pyarrow_item["media"][table_name] = lance_scanner.to_table()
  
          # Load PyArrow item from objects tables
          if load_objects:
              for table_name, obj_table in ds_tables["objects"].items():
                  lance_scanner = obj_table.to_lance().scanner(
                      filter=f"item_id in ('{item_id}')"
                  )
                  pyarrow_item["objects"][table_name] = lance_scanner.to_table()
  
          # Load PyArrow item from active learning tables
          if load_active_learning:
              for table_name, al_table in ds_tables["active_learning"].items():
                  lance_scanner = al_table.to_lance().scanner(
                      filter=f"id in ('{item_id}')"
                  )
                  pyarrow_item["active_learning"][table_name] = lance_scanner.to_table()
  
          # Load PyArrow item from segmentation embeddings tables
          found_embeddings = not load_embeddings
          if load_embeddings and "embeddings" in self.info.tables:
              for table in self.info.tables["embeddings"]:
                  if table.source.lower() in model_id.lower():
                      found_embeddings = True
                      emb_table = ds_tables["embeddings"][table.name]
                      lance_scanner = emb_table.to_lance().scanner(
                          filter=f"id in ('{item_id}')"
                      )
                      pyarrow_item["embeddings"][table.name] = lance_scanner.to_table()
  
          if pyarrow_item["main"]["db"].num_rows > 0 and found_embeddings:
              return DatasetItem.from_pyarrow(
                  pyarrow_item,
                  self.info,
                  self.media_dir,
                  media_features=True,
                  model_id=model_id,
              )
          return None
  
      def save_item(self, item: DatasetItem):
          """Save dataset item features and objects
  
          Args:
              item (DatasetItem): Item to save
          """
  
          # Update info in case of change
          self.info = self.load_info()
  
          # Load dataset tables
          ds_tables = self.open_tables()
          main_table = ds_tables["main"]["db"]
  
          # Add new item columns
          self.update_table(item, main_table, "main", "db")
  
          # Reload dataset tables
          ds_tables = self.open_tables()
          main_table = ds_tables["main"]["db"]
  
          # Update item
          item.update(main_table)
  
          # Add or update item objects
          for obj in item.objects.values():
              table_found = False
              if "objects" in self.info.tables:
                  for table in self.info.tables["objects"]:
                      if table.source == obj.source_id:
                          # Load object table
                          table_found = True
                          obj_table = ds_tables["objects"][table.name]
  
                          # Add new object columns
                          self.update_table(obj, obj_table, "objects", table.name)
  
                          # Reload dataset tables
                          ds_tables = self.open_tables()
                          main_table = ds_tables["objects"][table.name]
  
                          # Add or update object
                          obj.add_or_update(obj_table)
              # If first object
              if not table_found and obj.source_id == "Ground Truth":
                  # Create table
                  table = DatasetTable(
                      name="objects",
                      source="Ground Truth",
                      fields={
                          "id": "str",
                          "item_id": "str",
                          "view_id": "str",
                          "bbox": "bbox",
                          "mask": "compressedrle",
                      },
                  )
                  for feat in obj.features.values():
                      table.fields[feat.name] = feat.dtype
                  self.create_table(table, "objects")
  
                  # Reload dataset tables
                  ds_tables = self.open_tables()
                  obj_table = ds_tables["objects"][table.name]
  
                  # Add object
                  obj.add_or_update(obj_table)
  
          # Delete removed item objects
          item.delete_objects(ds_tables)
  
      def get_features_values(self, config_values: FeaturesValues) -> FeaturesValues:
          """get config values
             merge with distinct existing values for each scene and object string features, if not restricted
  
          Args:
              config_values (FeaturesValues): features values from db.json
          Returns:
              FeaturesValues: existing values for each scene and object string feature
          """
  
          # Load tables
          ds_tables = self.open_tables()
  
          def get_distinct_values(
              table_name: str,
              ignore_list: list[str],
              config_vals: dict[str, FeatureValues],
          ) -> dict[str, FeatureValues]:
              avail_values = defaultdict(FeatureValues)
              for table in ds_tables[table_name].values():
                  table_arrow = table.to_arrow()
                  feats = [f for f in table_arrow.column_names if f not in ignore_list]
                  for feat in feats:
                      avail_values[feat] = FeatureValues(restricted=False, values=[])
                      if config_vals and feat in config_vals:
                          avail_values[feat].restricted = config_vals[feat].restricted
                          avail_values[feat].values.extend(config_vals[feat].values)
                      if not avail_values[feat].restricted:
                          v = (
                              duckdb.sql(f"select DISTINCT {feat} from table_arrow")
                              .to_arrow_table()
                              .to_pydict()
                          )
                          if v[feat] is not None and v[feat] != [None]:
                              avail_values[feat].values.extend(
                                  [
                                      val
                                      for val in v[feat]
                                      if val is not None and isinstance(val, str)
                                  ]
                              )
              return avail_values
  
          return FeaturesValues(
              main=get_distinct_values(
                  "main",
                  ["id", "split", "views", "original_id"],
                  config_values.main if config_values else None,
              ),
              objects=get_distinct_values(
                  "objects",
                  ["id", "item_id", "view_id", "bbox", "mask", "review_state"],
                  config_values.objects if config_values else None,
              ),
          )
  
      @staticmethod
      def find(
          dataset_id: str,
          directory: Path | S3Path,
      ) -> "Dataset":
          """Find Dataset in directory
  
          Args:
              dataset_id (str): Dataset ID
              directory (Path): Directory to search in
  
          Returns:
              Dataset: Dataset
          """
  
          # Browse directory
          for json_fp in directory.glob("*/db.json"):
              info = DatasetInfo.from_json(json_fp)
              if info.id == dataset_id:
                  # Return dataset
                  return Dataset(json_fp.parent)
          return None
  
apis: lancedb.connect
extract_api: [((1770, 1810), 'pydantic.ConfigDict', 'ConfigDict', ([], {'arbitrary_types_allowed': '(True)'}), '(arbitrary_types_allowed=True)\n', (1780, 1810), False, 'from pydantic import BaseModel, ConfigDict\n'), ((5069, 5171), 'pixano.data.dataset.dataset_info.DatasetInfo.from_json', 'DatasetInfo.from_json', (["(self.path / 'db.json')"], {'load_stats': 'load_stats', 'load_thumbnail': 'load_thumbnail'}), "(self.path / 'db.json', load_stats=load_stats,\n    load_thumbnail=load_thumbnail)\n", (5090, 5171), False, 'from pixano.data.dataset.dataset_info import DatasetInfo\n'), ((6455, 6481), 'lancedb.connect', 'lancedb.connect', (['self.path'], {}), '(self.path)\n', (6470, 6481), False, 'import lancedb\n'), ((6798, 6815), 'collections.defaultdict', 'defaultdict', (['dict'], {}), '(dict)\n', (6809, 6815), False, 'from collections import defaultdict\n'), ((11462, 11479), 'collections.defaultdict', 'defaultdict', (['dict'], {}), '(dict)\n', (11473, 11479), False, 'from collections import defaultdict\n'), ((13728, 13745), 'collections.defaultdict', 'defaultdict', (['dict'], {}), '(dict)\n', (13739, 13745), False, 'from collections import defaultdict\n'), ((15029, 15046), 'collections.defaultdict', 'defaultdict', (['dict'], {}), '(dict)\n', (15040, 15046), False, 'from collections import defaultdict\n'), ((21782, 21799), 'collections.defaultdict', 'defaultdict', (['dict'], {}), '(dict)\n', (21793, 21799), False, 'from collections import defaultdict\n'), ((15990, 15996), 'pixano_inference.transformers.CLIP', 'CLIP', ([], {}), '()\n', (15994, 15996), False, 'from pixano_inference.transformers import CLIP\n'), ((23850, 23959), 'pixano.data.dataset.dataset_item.DatasetItem.from_pyarrow', 'DatasetItem.from_pyarrow', (['pyarrow_item', 'self.info', 'self.media_dir'], {'media_features': '(True)', 'model_id': 'model_id'}), '(pyarrow_item, self.info, self.media_dir,\n    media_features=True, model_id=model_id)\n', (23874, 23959), False, 'from pixano.data.dataset.dataset_item import DatasetItem\n'), ((27215, 27241), 'collections.defaultdict', 'defaultdict', (['FeatureValues'], {}), '(FeatureValues)\n', (27226, 27241), False, 'from collections import defaultdict\n'), ((29382, 29412), 'pixano.data.dataset.dataset_info.DatasetInfo.from_json', 'DatasetInfo.from_json', (['json_fp'], {}), '(json_fp)\n', (29403, 29412), False, 'from pixano.data.dataset.dataset_info import DatasetInfo\n'), ((2234, 2266), 'pixano.data.dataset.dataset_info.DatasetInfo.from_json', 'DatasetInfo.from_json', (['info_file'], {}), '(info_file)\n', (2255, 2266), False, 'from pixano.data.dataset.dataset_info import DatasetInfo\n'), ((9049, 9107), 'pixano.data.item.item_feature.ItemFeature', 'ItemFeature', ([], {'name': 'field_name', 'dtype': 'field_type', 'value': 'None'}), '(name=field_name, dtype=field_type, value=None)\n', (9060, 9107), False, 'from pixano.data.item.item_feature import FeaturesValues, FeatureValues, ItemFeature\n'), ((11751, 11870), 'duckdb.query', 'duckdb.query', (['f"""SELECT * FROM lance_table ORDER BY len({id_field}), {id_field} LIMIT {limit} OFFSET {offset}"""'], {}), "(\n    f'SELECT * FROM lance_table ORDER BY len({id_field}), {id_field} LIMIT {limit} OFFSET {offset}'\n    )\n", (11763, 11870), False, 'import duckdb\n'), ((12860, 12925), 'pixano.data.dataset.dataset_item.DatasetItem.from_pyarrow', 'DatasetItem.from_pyarrow', (['pyarrow_item', 'self.info', 'self.media_dir'], {}), '(pyarrow_item, self.info, self.media_dir)\n', (12884, 12925), False, 'from pixano.data.dataset.dataset_item import DatasetItem\n'), ((14187, 14252), 'pixano.data.dataset.dataset_item.DatasetItem.from_pyarrow', 'DatasetItem.from_pyarrow', (['pyarrow_item', 'self.info', 'self.media_dir'], {}), '(pyarrow_item, self.info, self.media_dir)\n', (14211, 14252), False, 'from pixano.data.dataset.dataset_item import DatasetItem\n'), ((17251, 17382), 'duckdb.query', 'duckdb.query', (['f"""SELECT id, _distance as distance FROM results_table ORDER BY _distance ASC LIMIT {limit} OFFSET {offset}"""'], {}), "(\n    f'SELECT id, _distance as distance FROM results_table ORDER BY _distance ASC LIMIT {limit} OFFSET {offset}'\n    )\n", (17263, 17382), False, 'import duckdb\n'), ((17538, 17645), 'duckdb.query', 'duckdb.query', (['"""SELECT * FROM results_table LEFT JOIN main_table USING (id) ORDER BY distance ASC"""'], {}), "(\n    'SELECT * FROM results_table LEFT JOIN main_table USING (id) ORDER BY distance ASC'\n    )\n", (17550, 17645), False, 'import duckdb\n'), ((18482, 18499), 'collections.defaultdict', 'defaultdict', (['dict'], {}), '(dict)\n', (18493, 18499), False, 'from collections import defaultdict\n'), ((25678, 25836), 'pixano.data.dataset.dataset_table.DatasetTable', 'DatasetTable', ([], {'name': '"""objects"""', 'source': '"""Ground Truth"""', 'fields': "{'id': 'str', 'item_id': 'str', 'view_id': 'str', 'bbox': 'bbox', 'mask':\n    'compressedrle'}"}), "(name='objects', source='Ground Truth', fields={'id': 'str',\n    'item_id': 'str', 'view_id': 'str', 'bbox': 'bbox', 'mask':\n    'compressedrle'})\n", (25690, 25836), False, 'from pixano.data.dataset.dataset_table import DatasetTable\n'), ((2286, 2319), 'pixano.data.dataset.dataset_stat.DatasetStat.from_json', 'DatasetStat.from_json', (['stats_file'], {}), '(stats_file)\n', (2307, 2319), False, 'from pixano.data.dataset.dataset_stat import DatasetStat\n'), ((12164, 12228), 'duckdb.query', 'duckdb.query', (['f"""SELECT * FROM lance_table WHERE id IN {id_list}"""'], {}), "(f'SELECT * FROM lance_table WHERE id IN {id_list}')\n", (12176, 12228), False, 'import duckdb\n'), ((27508, 27550), 'pixano.data.item.item_feature.FeatureValues', 'FeatureValues', ([], {'restricted': '(False)', 'values': '[]'}), '(restricted=False, values=[])\n', (27521, 27550), False, 'from pixano.data.item.item_feature import FeaturesValues, FeatureValues, ItemFeature\n'), ((3461, 3540), 'duckdb.query', 'duckdb.query', (['f"""SELECT id FROM lance_table WHERE original_id = \'{original_id}\'"""'], {}), '(f"SELECT id FROM lance_table WHERE original_id = \'{original_id}\'")\n', (3473, 3540), False, 'import duckdb\n'), ((4197, 4276), 'duckdb.query', 'duckdb.query', (['f"""SELECT id FROM lance_table WHERE original_id = \'{original_id}\'"""'], {}), '(f"SELECT id FROM lance_table WHERE original_id = \'{original_id}\'")\n', (4209, 4276), False, 'import duckdb\n'), ((7992, 8012), 'pixano.data.fields.Fields', 'Fields', (['table.fields'], {}), '(table.fields)\n', (7998, 8012), False, 'from pixano.data.fields import Fields, field_to_pyarrow\n'), ((10194, 10221), 'pixano.data.fields.field_to_pyarrow', 'field_to_pyarrow', (['col.dtype'], {}), '(col.dtype)\n', (10210, 10221), False, 'from pixano.data.fields import Fields, field_to_pyarrow\n'), ((10349, 10376), 'pixano.data.fields.field_to_pyarrow', 'field_to_pyarrow', (['col.dtype'], {}), '(col.dtype)\n', (10365, 10376), False, 'from pixano.data.fields import Fields, field_to_pyarrow\n'), ((12532, 12597), 'duckdb.query', 'duckdb.query', (['f"""SELECT * FROM lance_table  WHERE id IN {id_list}"""'], {}), "(f'SELECT * FROM lance_table  WHERE id IN {id_list}')\n", (12544, 12597), False, 'import duckdb\n'), ((16752, 16929), 'duckdb.query', 'duckdb.query', (['"""SELECT id, results_table._distance as distance_1, view_results_table._distance as distance_2 FROM results_table LEFT JOIN view_results_table USING (id)"""'], {}), "(\n    'SELECT id, results_table._distance as distance_1, view_results_table._distance as distance_2 FROM results_table LEFT JOIN view_results_table USING (id)'\n    )\n", (16764, 16929), False, 'import duckdb\n'), ((17008, 17147), 'duckdb.query', 'duckdb.query', (['"""SELECT (id), (SELECT Min(v) FROM (VALUES (distance_1), (distance_2)) AS value(v)) as _distance FROM results_table"""'], {}), "(\n    'SELECT (id), (SELECT Min(v) FROM (VALUES (distance_1), (distance_2)) AS value(v)) as _distance FROM results_table'\n    )\n", (17020, 17147), False, 'import duckdb\n'), ((19024, 19075), 'pyarrow.dataset.dataset', 'pa_ds.dataset', (["pyarrow_items['media'][media_source]"], {}), "(pyarrow_items['media'][media_source])\n", (19037, 19075), True, 'import pyarrow.dataset as pa_ds\n'), ((19116, 19133), 'pyarrow.dataset.field', 'pa_ds.field', (['"""id"""'], {}), "('id')\n", (19127, 19133), True, 'import pyarrow.dataset as pa_ds\n'), ((19943, 20001), 'pyarrow.dataset.dataset', 'pa_ds.dataset', (["pyarrow_items['active_learning'][al_source]"], {}), "(pyarrow_items['active_learning'][al_source])\n", (19956, 20001), True, 'import pyarrow.dataset as pa_ds\n'), ((27895, 27949), 'duckdb.sql', 'duckdb.sql', (['f"""select DISTINCT {feat} from table_arrow"""'], {}), "(f'select DISTINCT {feat} from table_arrow')\n", (27905, 27949), False, 'import duckdb\n'), ((20046, 20063), 'pyarrow.dataset.field', 'pa_ds.field', (['"""id"""'], {}), "('id')\n", (20057, 20063), True, 'import pyarrow.dataset as pa_ds\n')]