vastdb 0.0.5.2__py3-none-any.whl → 0.0.5.3__py3-none-any.whl

vastdb/api.py

@@ -757,6 +757,7 @@ class VastdbApi:
         if not port:
             port = 443 if secure else 80
 
+        self.default_max_list_columns_page_size = 1000
         self.session = requests.Session()
         self.session.verify = False
         self.session.headers['user-agent'] = "VastData Tabular API 1.0 - 2022 (c)"
@@ -1207,7 +1208,7 @@ class VastdbApi:
                                 data=serialized_schema, headers=headers)
         return self._check_res(res, "drop_columns", expected_retvals)
 
-    def list_columns(self, bucket, schema, table, *, txid=0, client_tags=None, max_keys=1000, next_key=0,
+    def list_columns(self, bucket, schema, table, *, txid=0, client_tags=None, max_keys=None, next_key=0,
                      count_only=False, name_prefix="", exact_match=False,
                      expected_retvals=None, bc_list_internals=False):
         """
@@ -1218,6 +1219,7 @@ class VastdbApi:
         tabular-max-keys: 1000
         tabular-next-key: NextColumnId
         """
+        max_keys = max_keys or self.default_max_list_columns_page_size
         client_tags = client_tags or []
         expected_retvals = expected_retvals or []
 
@@ -2007,7 +2009,7 @@ class VastdbApi:
         txid, created_txid = self._begin_tx_if_necessary(txid)
 
         if rows:
-            columns = self._list_table_columns(bucket, schema, table, field_names=rows.keys())
+            columns = self._list_table_columns(bucket, schema, table, field_names=rows.keys(), txid=txid)
             columns_dict = dict([(column[0], column[1]) for column in columns])
             arrow_schema = pa.schema([])
             arrays = []

vastdb/bench_scan.py

@@ -0,0 +1,45 @@
+from vastdb import api
+
+from logbook import Logger, StreamHandler
+import sys
+import time
+import pprint
+
+StreamHandler(sys.stdout).push_application()
+log = Logger('Logbook')
+
+# access_key_id=F3YUMQZDQB60ZZJ1PBAZ
+# secret_access_key=9a9Q3if6IC5LjUexly/nXFv1UCANBnhGxi++Sw6p
+
+a = api.VastdbApi(
+    access_key='F3YUMQZDQB60ZZJ1PBAZ', 
+    secret_key='9a9Q3if6IC5LjUexly/nXFv1UCANBnhGxi++Sw6p', 
+    host='172.19.111.1:172.19.111.16')
+
+kwargs = dict(
+    bucket='tabular-slothful-jocular-jack',
+    schema='tpcds_schema_create_as_select',
+    table='store_sales',
+    field_names=['ss_sold_date_sk', 'ss_sold_time_sk', 'ss_item_sk'],
+    filters={'ss_item_sk': ['le 1']},
+    num_sub_splits=8)
+
+pprint.pprint(kwargs)
+
+res = a.query_iterator(**kwargs)
+
+total_bytes = 0
+total_rows = 0
+start = time.time()
+last_log = None
+
+for b in res: 
+    total_bytes += b.get_total_buffer_size()
+    total_rows += len(b)
+    dt = time.time() - start
+    if last_log != int(dt):
+        log.info("{:.3f} Mrow/s, {:.3f} MB/s", (total_rows/dt) / 1e6, (total_bytes/dt) / 1e6)
+        last_log = int(dt)
+
+dt = time.time() - start
+log.info("Done after {:.3f} seconds, {:.3f} Mrows, {:.3f} MB", dt, total_rows / 1e6, total_bytes / 1e6)

vastdb/tests/conftest.py

@@ -0,0 +1,45 @@
+import pytest
+import boto3
+
+from vastdb import v2
+
+
+def pytest_addoption(parser):
+    parser.addoption("--tabular-bucket-name", help="Name of the S3 bucket with Tabular enabled")
+    parser.addoption("--tabular-access-key", help="Access key with Tabular permissions")
+    parser.addoption("--tabular-secret-key", help="Secret key with Tabular permissions")
+    parser.addoption("--tabular-endpoint-url", help="Tabular server endpoint")
+
+
+@pytest.fixture(scope="module")
+def rpc(request):
+    return v2.connect(
+        access=request.config.getoption("--tabular-access-key"),
+        secret=request.config.getoption("--tabular-secret-key"),
+        endpoint=request.config.getoption("--tabular-endpoint-url"),
+    )
+
+
+@pytest.fixture(scope="module")
+def test_bucket_name(request):
+    return request.config.getoption("--tabular-bucket-name")
+
+
+@pytest.fixture(scope="module")
+def clean_bucket_name(request, test_bucket_name, rpc):
+    with rpc.transaction() as tx:
+        b = tx.bucket(test_bucket_name)
+        for s in b.schemas():
+            for t in s.tables():
+                t.drop()
+            s.drop()
+    return test_bucket_name
+
+
+@pytest.fixture(scope="module")
+def s3(request):
+    return boto3.client(
+        's3',
+        aws_access_key_id=request.config.getoption("--tabular-access-key"),
+        aws_secret_access_key=request.config.getoption("--tabular-secret-key"),
+        endpoint_url=request.config.getoption("--tabular-endpoint-url"))

vastdb/tests/test_create_table_from_parquets.py

@@ -0,0 +1,50 @@
+import pytest
+import os
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from vastdb.v2 import InvalidArgumentError
+from vastdb import util
+
+
+def test_create_table_from_files(rpc, clean_bucket_name, s3):
+    datasets = [
+        {'num': [0],
+         'varch': ['z']},
+        {'num': [1, 2, 3, 4, 5],
+         'varch': ['a', 'b', 'c', 'd', 'e']},
+        {'num': [1, 2, 3, 4, 5],
+         'bool': [True, False, None, None, False],
+         'varch': ['a', 'b', 'c', 'd', 'e']},
+        {'num': [1, 2],
+         'bool': [True, True]},
+        {'varch': ['a', 'b', 'c'],
+         'mismatch': [1, 2, 3]}
+    ]
+    for i, ds in enumerate(datasets):
+        table = pa.Table.from_pydict(ds)
+        pq.write_table(table, f'prq{i}')
+        with open(f'prq{i}', 'rb') as f:
+            s3.put_object(Bucket=clean_bucket_name, Key=f'prq{i}', Body=f)
+        os.remove(f'prq{i}')
+
+    same_schema_files = [f'/{clean_bucket_name}/prq{i}' for i in range(2)]
+    contained_schema_files = [f'/{clean_bucket_name}/prq{i}' for i in range(4)]
+    different_schema_files = [f'/{clean_bucket_name}/prq{i}' for i in range(5)]
+
+    with rpc.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        s = b.create_schema('s1')
+        t = util.create_table_from_files(s, 't1', contained_schema_files)
+        assert len(t.arrow_schema) == 3
+        assert t.arrow_schema == pa.schema([('num', pa.int64()), ('bool', pa.bool_()), ('varch', pa.string())])
+
+        with pytest.raises(InvalidArgumentError):
+            util.create_table_from_files(s, 't2', different_schema_files)
+
+        with pytest.raises(InvalidArgumentError):
+            util.create_table_from_files(s, 't2', contained_schema_files, schema_merge_func=util.strict_schema_merge)
+
+        util.create_table_from_files(s, 't2', different_schema_files, schema_merge_func=util.union_schema_merge)
+        util.create_table_from_files(s, 't3', same_schema_files, schema_merge_func=util.strict_schema_merge)

vastdb/tests/test_sanity.py

@@ -0,0 +1,63 @@
+import logging
+
+import threading
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from vastdb import api
+from itertools import cycle
+
+log = logging.getLogger(__name__)
+
+def test_hello_world(rpc):
+    with rpc.transaction() as tx:
+        assert tx.txid is not None
+
+def test_version_extraction():
+    # A list of version and expected version parsed by API
+    TEST_CASES = [
+            (None, None), # vast server without version in header
+            ("5", None),                                    # major only is not supported
+            ("5.2", "5.2"),                                 # major.minor
+            ("5.2.0", "5.2.0"),                             # major.minor.patch
+            ("5.2.0.0", "5.2.0.0"),                         # major.minor.patch.protocol
+            ("5.2.0.0 some other things", "5.2.0.0"),       # Test forward comptibility 1
+            ("5.2.0.0.20 some other things", "5.2.0.0"),    # Test forward comptibility 2
+    ]
+
+    # Mock OPTIONS handle that cycles through the test cases response
+    class MockOptionsHandler(BaseHTTPRequestHandler):
+        versions_iterator = cycle(TEST_CASES)
+
+        def __init__(self, *args) -> None:
+            super().__init__(*args)
+
+        def do_OPTIONS(self):
+            self.send_response(204)
+            self.end_headers()
+
+        def version_string(self):
+            version = next(self.versions_iterator)[0]
+            return f"vast {version}" if version else "vast"
+
+        def log_message(self, format, *args):
+            log.debug(format,*args)
+
+    # start the server on localhost on some available port port
+    server_address =('localhost', 0)
+    httpd = HTTPServer(server_address, MockOptionsHandler)
+
+    def start_http_server_in_thread():
+        log.info(f"Mock HTTP server is running on port {httpd.server_port}")
+        httpd.serve_forever()
+        log.info("Mock HTTP server killed")
+
+    # start the server in a thread so we have the main thread to operate the API
+    server_thread = threading.Thread(target=start_http_server_in_thread)
+    server_thread.start()
+
+    try:
+        for test_case in TEST_CASES:
+            tester = api.VastdbApi(endpoint=f"http://localhost:{httpd.server_port}", access_key="abc", secret_key="abc")
+            assert tester.vast_version == test_case[1]
+    finally:
+        # make sure we shut the server down no matter what
+        httpd.shutdown()

vastdb/tests/test_schemas.py

@@ -0,0 +1,39 @@
+import pytest
+
+
+def test_schemas(rpc, clean_bucket_name):
+    with rpc.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        assert b.schemas() == []
+
+        s = b.create_schema('s1')
+        assert s.bucket == b
+        assert b.schemas() == [s]
+
+        s.rename('s2')
+        assert s.bucket == b
+        assert s.name == 's2'
+        assert b.schemas()[0].name == 's2'
+
+        s.drop()
+        assert b.schemas() == []
+
+
+def test_commits_and_rollbacks(rpc, clean_bucket_name):
+    with rpc.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        assert b.schemas() == []
+        b.create_schema("s3")
+        assert b.schemas() != []
+        # implicit commit
+
+    with pytest.raises(ZeroDivisionError):
+        with rpc.transaction() as tx:
+            b = tx.bucket(clean_bucket_name)
+            b.schema("s3").drop()
+            assert b.schemas() == []
+            1/0  # rollback schema dropping
+
+    with rpc.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        assert b.schemas() != []

vastdb/tests/test_tables.py

@@ -0,0 +1,40 @@
+import pyarrow as pa
+
+
+def test_tables(rpc, clean_bucket_name):
+    with rpc.transaction() as tx:
+        s = tx.bucket(clean_bucket_name).create_schema('s1')
+        columns = pa.schema([
+            ('a', pa.int16()),
+            ('b', pa.float32()),
+            ('s', pa.utf8()),
+        ])
+        assert s.tables() == []
+        t = s.create_table('t1', columns)
+        assert s.tables() == [t]
+
+        rb = pa.record_batch(schema=columns, data=[
+            [111, 222],
+            [0.5, 1.5],
+            ['a', 'b'],
+        ])
+        expected = pa.Table.from_batches([rb])
+        t.insert(rb)
+
+        actual = pa.Table.from_batches(t.select(columns=['a', 'b', 's']))
+        assert actual == expected
+
+        actual = pa.Table.from_batches(t.select(columns=['a', 'b']))
+        assert actual == expected.select(['a', 'b'])
+
+        actual = pa.Table.from_batches(t.select(columns=['b', 's', 'a']))
+        assert actual == expected.select(['b', 's', 'a'])
+
+        actual = pa.Table.from_batches(t.select(columns=['s']))
+        assert actual == expected.select(['s'])
+
+        actual = pa.Table.from_batches(t.select(columns=[]))
+        assert actual == expected.select([])
+
+        t.drop()
+        s.drop()

vastdb/util.py

@@ -0,0 +1,77 @@
+import logging
+from typing import Callable
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from vastdb.v2 import InvalidArgumentError, Table, Schema
+
+
+log = logging.getLogger(__name__)
+
+
+def create_table_from_files(
+        schema: Schema, table_name: str, parquet_files: [str], schema_merge_func: Callable = None) -> Table:
+    if not schema_merge_func:
+        schema_merge_func = default_schema_merge
+    else:
+        assert schema_merge_func in [default_schema_merge, strict_schema_merge, union_schema_merge]
+    tx = schema.tx
+    current_schema = pa.schema([])
+    s3fs = pa.fs.S3FileSystem(
+        access_key=tx._rpc.api.access_key, secret_key=tx._rpc.api.secret_key, endpoint_override=tx._rpc.api.url)
+    for prq_file in parquet_files:
+        if not prq_file.startswith('/'):
+            raise InvalidArgumentError(f"Path {prq_file} must start with a '/'")
+        parquet_ds = pq.ParquetDataset(prq_file.lstrip('/'), filesystem=s3fs)
+        current_schema = schema_merge_func(current_schema, parquet_ds.schema)
+
+
+    log.info("Creating table %s from %d Parquet files, with columns: %s",
+             table_name, len(parquet_files), list(current_schema))
+    table = schema.create_table(table_name, current_schema)
+
+    log.info("Starting import of %d files to table: %s", len(parquet_files), table)
+    table.import_files(parquet_files)
+    log.info("Finished import of %d files to table: %s", len(parquet_files), table)
+    return table
+
+
+def default_schema_merge(current_schema: pa.Schema, new_schema: pa.Schema) -> pa.Schema:
+    """
+    This function validates a schema is contained in another schema
+    Raises an InvalidArgumentError if a certain field does not exist in the target schema
+    """
+    if not current_schema.names:
+        return new_schema
+    s1 = set(current_schema)
+    s2 = set(new_schema)
+
+    if len(s1) > len(s2):
+        s1, s2 = s2, s1
+        result = current_schema  # We need this variable in order to preserve the original fields order
+    else:
+        result = new_schema
+
+    if not s1.issubset(s2):
+        log.error("Schema mismatch. schema: %s isn't contained in schema: %s.", s1, s2)
+        raise InvalidArgumentError("Found mismatch in parquet files schemas.")
+    return result
+
+
+def strict_schema_merge(current_schema: pa.Schema, new_schema: pa.Schema) -> pa.Schema:
+    """
+    This function validates two Schemas are identical.
+    Raises an InvalidArgumentError if schemas aren't identical.
+    """
+    if current_schema.names and current_schema != new_schema:
+        raise InvalidArgumentError(f"Schemas are not identical. \n {current_schema} \n vs \n {new_schema}")
+
+    return new_schema
+
+
+def union_schema_merge(current_schema: pa.Schema, new_schema: pa.Schema) -> pa.Schema:
+    """
+    This function returns a unified schema from potentially two different schemas.
+    """
+    return pa.unify_schemas([current_schema, new_schema])

vastdb/v2.py

@@ -1,108 +1,360 @@
-from vastdb import *
+from dataclasses import dataclass, field
+import logging
+import os
 
+import boto3
+import botocore
+import ibis
+import pyarrow as pa
+import requests
 
-class Context:
-	tx: int
-	_rpc: RPC
+from vastdb.api import VastdbApi, serialize_record_batch, build_query_data_request, parse_query_data_response, TABULAR_INVALID_ROW_ID
 
-	def bucket(name: str) -> Bucket
 
+log = logging.getLogger(__name__)
+
+
+class VastException(Exception):
+    pass
+
+
+class NotFoundError(VastException):
+    pass
+
+
+class AccessDeniedError(VastException):
+    pass
+
+
+class ImportFilesError(VastException):
+    pass
+
+
+class InvalidArgumentError(VastException):
+    pass
+
+
+class RPC:
+    def __init__(self, access=None, secret=None, endpoint=None):
+        if access is None:
+            access = os.environ['AWS_ACCESS_KEY_ID']
+        if secret is None:
+            secret = os.environ['AWS_SECRET_ACCESS_KEY']
+        if endpoint is None:
+            endpoint = os.environ['AWS_S3_ENDPOINT_URL']
+
+        self.api = VastdbApi(endpoint, access, secret)
+        self.s3 = boto3.client('s3',
+            aws_access_key_id=access,
+            aws_secret_access_key=secret,
+            endpoint_url=endpoint)
+
+    def __repr__(self):
+        return f'RPC(endpoint={self.api.url}, access={self.api.access_key})'
+
+    def transaction(self):
+        return Transaction(self)
+
+
+def connect(*args, **kw):
+    return RPC(*args, **kw)
+
+
+@dataclass
+class Transaction:
+    _rpc: RPC
+    txid: int = None
+
+    def __enter__(self):
+        response = self._rpc.api.begin_transaction()
+        self.txid = int(response.headers['tabular-txid'])
+        log.debug("opened txid=%016x", self.txid)
+        return self
+
+    def __exit__(self, *args):
+        if args == (None, None, None):
+            log.debug("committing txid=%016x", self.txid)
+            self._rpc.api.commit_transaction(self.txid)
+        else:
+            log.exception("rolling back txid=%016x", self.txid)
+            self._rpc.api.rollback_transaction(self.txid)
+
+    def __repr__(self):
+        return f'Transaction(id=0x{self.txid:016x})'
+
+    def bucket(self, name: str) -> "Bucket":
+        try:
+            self._rpc.s3.head_bucket(Bucket=name)
+            return Bucket(name, self)
+        except botocore.exceptions.ClientError as e:
+            if e.response['Error']['Code'] == 403:
+                raise AccessDeniedError(f"Access is denied to bucket: {name}") from e
+            else:
+                raise NotFoundError(f"Bucket {name} does not exist") from e
+
+
+@dataclass
 class Bucket:
-	ctx: Context
-	name: str
+    name: str
+    tx: Transaction
+
+    def create_schema(self, path: str) -> "Schema":
+        self.tx._rpc.api.create_schema(self.name, path, txid=self.tx.txid)
+        log.info("Created schema: %s", path)
+        return self.schema(path)
 
-	def schema(name: str) -> Schema
+    def schema(self, path: str) -> "Schema":
+        schema = self.schemas(path)
+        log.debug("schema: %s", schema)
+        if not schema:
+            raise NotFoundError(f"Schema '{path}' was not found in bucket: {self.name}")
+        assert len(schema) == 1, f"Expected to receive only a single schema, but got: {len(schema)}. ({schema})"
+        log.debug("Found schema: %s", schema[0].name)
+        return schema[0]
 
+    def schemas(self, schema: str = None) -> ["Schema"]:
+        schemas = []
+        next_key = 0
+        exact_match = bool(schema)
+        log.debug("list schemas param: schema=%s, exact_match=%s", schema, exact_match)
+        while True:
+            bucket_name, curr_schemas, next_key, is_truncated, _ = \
+                self.tx._rpc.api.list_schemas(bucket=self.name, next_key=next_key, txid=self.tx.txid,
+                                               name_prefix=schema, exact_match=exact_match)
+            if not curr_schemas:
+                break
+            schemas.extend(curr_schemas)
+            if not is_truncated:
+                break
+
+        return [Schema(name=name, bucket=self) for name, *_ in schemas]
+
+
+@dataclass
 class Schema:
-	ctx: Context
-	path: str
+    name: str
+    bucket: Bucket
+
+    @property
+    def tx(self):
+        return self.bucket.tx
+
+    def create_table(self, table_name: str, columns: pa.Schema) -> "Table":
+        self.tx._rpc.api.create_table(self.bucket.name, self.name, table_name, columns, txid=self.tx.txid)
+        log.info("Created table: %s", table_name)
+        return self.table(table_name)
+
+    def table(self, name: str) -> "Table":
+        t = self.tables(table_name=name)
+        if not t:
+            raise NotFoundError(f"Table '{name}' was not found under schema: {self.name}")
+        assert len(t) == 1, f"Expected to receive only a single table, but got: {len(t)}. tables: {t}"
+        log.debug("Found table: %s", t[0])
+        return t[0]
+
+    def tables(self, table_name=None) -> ["Table"]:
+        tables = []
+        next_key = 0
+        name_prefix = table_name if table_name else ""
+        exact_match = bool(table_name)
+        while True:
+            bucket_name, schema_name, curr_tables, next_key, is_truncated, _ = \
+                self.tx._rpc.api.list_tables(
+                    bucket=self.bucket.name, schema=self.name, next_key=next_key, txid=self.tx.txid,
+                    exact_match=exact_match, name_prefix=name_prefix)
+            if not curr_tables:
+                break
+            tables.extend(curr_tables)
+            if not is_truncated:
+                break
+
+        return [_parse_table_info(table, self) for table in tables]
+
+    def drop(self) -> None:
+        self.tx._rpc.api.drop_schema(self.bucket.name, self.name, txid=self.tx.txid)
+        log.info("Dropped schema: %s", self.name)
+
+    def rename(self, new_name) -> None:
+        self.tx._rpc.api.alter_schema(self.bucket.name, self.name, txid=self.tx.txid, new_name=new_name)
+        log.info("Renamed schema: %s to %s", self.name, new_name)
+        self.name = new_name
+
+
+@dataclass
+class TableStats:
+    num_rows: int
+    size: int
+
+
+@dataclass
+class QueryConfig:
+    num_sub_splits: int = 4
+    num_splits: int = 1
+    data_endpoints: [str] = None
+    limit_per_sub_split: int = 128 * 1024
+    num_row_groups_per_sub_split: int = 8
 
-	def schema(name: str) -> Schema
-	def table(name: str) -> Table
 
+@dataclass
 class Table:
-	ctx: Context
-	path: str
+    name: str
+    schema: pa.Schema
+    handle: int
+    stats: TableStats
+    properties: dict = None
+    arrow_schema: pa.Schema = field(init=False, compare=False)
+    _ibis_table: ibis.Schema = field(init=False, compare=False)
 
-	def import_files(...)
-	def import_partitioned_files(...)
-	def select(...) -> ???
+    def __post_init__(self):
+        self.properties = self.properties or {}
+        self.arrow_schema = self.columns()
+        self._ibis_table = ibis.Schema.from_pyarrow(self.arrow_schema)
 
+    @property
+    def tx(self):
+        return self.schema.tx
 
-class RPC:
-	"""
-	INTERNAL STUFF: actually uses requests to send/receive stuff
-	Cannot do pagination
-	"""
+    @property
+    def bucket(self):
+        return self.schema.bucket
 
-	### We can just copy-paste stuff from api.py
+    def __repr__(self):
+        return f"{type(self).__name__}(name={self.name})"
 
-	def single_shot_query_data()
-	def single_shot_list_columns()
+    def columns(self) -> pa.Schema:
+        cols = self.tx._rpc.api._list_table_columns(self.bucket.name, self.schema.name, self.name, txid=self.tx.txid)
+        self.arrow_schema = pa.schema([(col[0], col[1]) for col in cols])
+        return self.arrow_schema
 
-@contextmanager
-def context(access, secret, endpoint):
-	rpc = RPC(access, secret, endpoint) # Low-level commands => the user should not use it
-	tx = rpc.begin_transaction()
-	try:
-		yield Context(rpc, tx)
-	finally:
-		rpc.close_transaction(tx)
+    def import_files(self, files_to_import: [str]) -> None:
+        source_files = {}
+        for f in files_to_import:
+            bucket_name, object_path = _parse_bucket_and_object_names(f)
+            source_files[(bucket_name, object_path)] = b''
 
+        self._execute_import(source_files)
 
-with context(access, secret, endpoint) as ctx:  # open/closes tx
-	# tx keep-alive?
-	b = ctx.bucket("buck") # may raise NotFoundError if bucket is missing
+    def import_partitioned_files(self, files_and_partitions: {str: pa.RecordBatch}) -> None:
+        source_files = {}
+        for f, record_batch in files_and_partitions.items():
+            bucket_name, object_path = _parse_bucket_and_object_names(f)
+            serialized_batch = _serialize_record_batch(record_batch)
+            source_files = {(bucket_name, object_path): serialized_batch.to_pybytes()}
 
-	ctx._rpc.strange_thing???
+        self._execute_import(source_files)
 
-	b.create_schema("s1")
-	b.create_schema("s1/s2")
-	b.create_schema("s1/s2/s3")
+    def _execute_import(self, source_files):
+        try:
+            self.tx._rpc.api.import_data(
+                self.bucket.name, self.schema.name, self.name, source_files, txid=self.tx.txid)
+        except requests.HTTPError as e:
+            raise ImportFilesError(f"import_files failed with status: {e.response.status_code}, reason: {e.response.reason}")
+        except Exception as e:
+            # TODO: investigate and raise proper error in case of failure mid import.
+            raise ImportFilesError("import_files failed") from e
 
-	iterable_of_schema_objects = b.schemas() # BFS or only top-level?
+    def select(self, columns: [str], predicate: ibis.expr.types.BooleanColumn = None,
+               config: "QueryConfig" = None):
+        if config is None:
+            config = QueryConfig()
 
-	s = b.schema("s1/s2/s3") # may raise NotFoundError if schema is missing
-	s = b.schema("s1").schema("s2/s3") # may raise NotFoundError if schema is missing
-	s = b / "s1" / "s2" / "s3" # may raise NotFoundError if schema is missing
+        api = self.tx._rpc.api
+        field_names = columns
+        filters = []
+        bucket = self.bucket.name
+        schema = self.schema.name
+        table = self.name
+        query_data_request = build_query_data_request(
+            schema=self.arrow_schema, filters=filters, field_names=field_names)
 
-	assert s.schemas() == []
+        start_row_ids = {i: 0 for i in range(config.num_sub_splits)}
+        assert config.num_splits == 1  # TODO()
+        split = (0, 1, config.num_row_groups_per_sub_split)
+        response_row_id = False
 
-	iterable_of_tables_objects = s.tables()
-	t = s.table("t") # /bucket/s1/s2/s3/t under tx
+        while not all(row_id == TABULAR_INVALID_ROW_ID for row_id in start_row_ids.values()):
+            response = api.query_data(
+                bucket=bucket,
+                schema=schema,
+                table=table,
+                params=query_data_request.serialized,
+                split=split,
+                num_sub_splits=config.num_sub_splits,
+                response_row_id=response_row_id,
+                txid=self.tx.txid,
+                limit_rows=config.limit_per_sub_split,
+                sub_split_start_row_ids=start_row_ids.items())
 
-	s.rename()
-	s.drop()
-	...
+            pages_iter = parse_query_data_response(
+                conn=response.raw,
+                schema=query_data_request.response_schema,
+                start_row_ids=start_row_ids)
 
+            for page in pages_iter:
+                for batch in page.to_batches():
+                    if len(batch) > 0:
+                        yield batch
 
-	# may take a while - finishes when all files are done
-	# if all OK, return None
-	# in case of error raise ImportFilesError(failed_files_list=[(path, code, reason)])
-	t.import_files(["/buck1/file1", ... "/buck3/file3"])
-	t.import_partitioned_files({"/buck1/file1": pa.RecordBatch, ... "/buck3/file3": pa.RecordBatch})
+    def insert(self, rows: pa.RecordBatch) -> None:
+        blob = serialize_record_batch(rows)
+        self.tx._rpc.api.insert_rows(self.bucket.name, self.schema.name, self.name, record_batch=blob, txid=self.tx.txid)
 
-	arrow_schema = t.columns()
-	iterable_of_record_batches = t.select(
-		column_names: List[str], 
-		predicate: ibis.BooleanColumn???,
-		limit: int = None,
-		config: QueryConfig = None
-	)
+    def drop(self) -> None:
+        self.tx._rpc.api.drop_table(self.bucket.name, self.schema.name, self.name, txid=self.tx.txid)
+        log.info("Dropped table: %s", self.name)
 
+    def rename(self, new_name) -> None:
+        self.tx._rpc.api.alter_table(
+            self.bucket.name, self.schema.name, self.name, txid=self.tx.txid, new_name=new_name)
+        log.info("Renamed table from %s to %s ", self.name, new_name)
+        self.name = new_name
 
-	t.drop()
-	t.rename()
-	t.add_column()
-	t.drop_column()
-	...
+    def add_column(self, new_column: pa.Schema) -> None:
+        self.tx._rpc.api.add_columns(self.bucket.name, self.schema.name, self.name, new_column, txid=self.tx.txid)
+        log.info("Added column(s): %s", new_column)
+        self.arrow_schema = self.columns()
 
+    def drop_column(self, column_to_drop: pa.Schema) -> None:
+        self.tx._rpc.api.drop_columns(self.bucket.name, self.schema.name, self.name, column_to_drop, txid=self.tx.txid)
+        log.info("Dropped column(s): %s", column_to_drop)
+        self.arrow_schema = self.columns()
 
-class QueryConfig:
-	num_of_subsplits: int = 2
-	num_of_splits: int = 16?
-	# how to load balance between VIPs?
-	# we need a new RPC to get the "data_enpoints" VIPs from VAST and then we can round-robin between them?
-	# => @alon
-	##### list_of_data_endpoints: List[str] = None
-	limit_per_sub_split: int = 128k
+    def rename_column(self, current_column_name: str, new_column_name: str) -> None:
+        self.tx._rpc.api.alter_column(self.bucket.name, self.schema.name, self.name, name=current_column_name,
+                                       new_name=new_column_name, txid=self.tx.txid)
+        log.info("Renamed column: %s to %s", current_column_name, new_column_name)
+        self.arrow_schema = self.columns()
+
+    def __getitem__(self, col_name):
+        return self._ibis_table[col_name]
+
+
+def _parse_table_info(table_info, schema: "Schema"):
+    stats = TableStats(num_rows=table_info.num_rows, size=table_info.size_in_bytes)
+    return Table(name=table_info.name, schema=schema, handle=int(table_info.handle), stats=stats)
+
+
+def _parse_bucket_and_object_names(path: str) -> (str, str):
+    if not path.startswith('/'):
+        raise InvalidArgumentError(f"Path {path} must start with a '/'")
+    components = path.split(os.path.sep)
+    bucket_name = components[1]
+    object_path = os.path.sep.join(components[2:])
+    return bucket_name, object_path
+
+
+def _serialize_record_batch(record_batch: pa.RecordBatch) -> pa.lib.Buffer:
+    sink = pa.BufferOutputStream()
+    with pa.ipc.new_stream(sink, record_batch.schema) as writer:
+        writer.write(record_batch)
+    return sink.getvalue()
+
+
+def _parse_endpoint(endpoint):
+    if ":" in endpoint:
+        endpoint, port = endpoint.split(":")
+        port = int(port)
+    else:
+        port = 80
+    log.debug("endpoint: %s, port: %d", endpoint, port)
+    return endpoint, port

{vastdb-0.0.5.2.dist-info → vastdb-0.0.5.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vastdb
-Version: 0.0.5.2
+Version: 0.0.5.3
 Summary: VAST Data SDK
 Home-page: https://github.com/vast-data/vastdb_sdk
 Author: VAST DATA

{vastdb-0.0.5.2.dist-info → vastdb-0.0.5.3.dist-info}/RECORD

@@ -163,10 +163,18 @@ vast_protobuf/substrait/extensions/extensions_pb2.py,sha256=I_6c6nMmMaYvVtzF-5yc
 vast_protobuf/tabular/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vast_protobuf/tabular/rpc_pb2.py,sha256=7kW2WrA2sGk6WVbD83mc_cKkZ2MxoImSO5GOVz6NbbE,23776
 vastdb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vastdb/api.py,sha256=ZfkwDyTqaQ_j1wmp2iog9Yj3LW98S8IzLxs7gExiPG8,124746
-vastdb/v2.py,sha256=0fLulaIQGlIbVNBBFGd6iwYPuGhaaJIHTiJORyio_YQ,2438
-vastdb-0.0.5.2.dist-info/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
-vastdb-0.0.5.2.dist-info/METADATA,sha256=t6yxMkCPHuy7GJU6lSmd6QjrCsGmcPvzJE0qihLjMQQ,1369
-vastdb-0.0.5.2.dist-info/WHEEL,sha256=ewwEueio1C2XeHTvT17n8dZUJgOvyCWCt0WVNLClP9o,92
-vastdb-0.0.5.2.dist-info/top_level.txt,sha256=Vsj2MKtlhPg0J4so64slQtnwjhgoPmJgcG-6YcVAwVc,20
-vastdb-0.0.5.2.dist-info/RECORD,,
+vastdb/api.py,sha256=u5Cf01LeHGN7x_pcjnzfLV-lU485FGFCv7eTIKpSaB0,124883
+vastdb/bench_scan.py,sha256=95O34oHS0UehX2ad4T2mok87CKszCFLCDZASMnZp77M,1208
+vastdb/util.py,sha256=EF892Gbs08BxHVgG3FZ6QvhpKI2-eIL5bPzzrYE_Qd8,2905
+vastdb/v2.py,sha256=gWZUnhSLEvtrXPxoTpTAwNuzU9qxrCaWKXmeNBpMrGY,12601
+vastdb/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vastdb/tests/conftest.py,sha256=dcYFJO0Riyn687qZTwcwKbwGieg6s4yZrVFrJAX-ylU,1461
+vastdb/tests/test_create_table_from_parquets.py,sha256=dxykmvUR-vui6Z3qUvXPYJ9Nw6V_qcxKl4NDNQK4kiY,1963
+vastdb/tests/test_sanity.py,sha256=7HmCjuOmtoYnuWiPjMP6m7sYQYop1_qRCzq2ZX0rKlc,2404
+vastdb/tests/test_schemas.py,sha256=-nntn3ltBaaqSTsUvi-i9J0yr4TYvOTRyTNY039vEIk,1047
+vastdb/tests/test_tables.py,sha256=KPe0ESVGWixecTSwQ8whzSF-NZrNVZ-Kv-C4Gz-OQnQ,1225
+vastdb-0.0.5.3.dist-info/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
+vastdb-0.0.5.3.dist-info/METADATA,sha256=Yd93AoZE5ZUhJUr0MhtfhcMaQUtSFZ1wbzc6vvEvclQ,1369
+vastdb-0.0.5.3.dist-info/WHEEL,sha256=ewwEueio1C2XeHTvT17n8dZUJgOvyCWCt0WVNLClP9o,92
+vastdb-0.0.5.3.dist-info/top_level.txt,sha256=Vsj2MKtlhPg0J4so64slQtnwjhgoPmJgcG-6YcVAwVc,20
+vastdb-0.0.5.3.dist-info/RECORD,,