#  Copyright (c) 2024. Lance Developers
#
#  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.

import pyarrow as pa

def train_ivf_model(
    dataset,
    column: str,
    dimension: int,
    num_partitions: int,
    distance_type: str,
    sample_rate: int,
    max_iters: int,
) -> pa.Array: ...
def train_pq_model(
    dataset,
    column: str,
    dimension: int,
    num_subvectors: int,
    distance_type: str,
    sample_rate: int,
    max_iters: int,
    ivf_model: pa.Array,
) -> pa.Array: ...
def transform_vectors(
    dataset,
    column: str,
    dimension: int,
    num_subvectors: int,
    distance_type: str,
    ivf_centroids: pa.Array,
    pq_codebook: pa.Array,
    dst_uri: str,
): ...